Ways to calculate diagonal Hessian

[shuhuah]

Dear Alan,

Thank you very much for publicly sharing deepwave. It is a very nice tool for testing different ideas in FWI.

I have one question regarding the Hessian calculation. You shared nice example of two times backpropagation to calculate full Hessian, torch.autograd.functional.hessian(wrap, v). However, in many cases it is intractable due to computation and memory limitations. Now I am trying to calculate only the diagonal elements of Hessian using backpropagation.

I am trying to use torch.autograd.grad() for this purpose. First I calculate the first order derivative using gradient = torch.autograd.grad(loss, v, create_graph=True)[0], which works fine. Next I want to calculate the gradient of gradient for each single velocity parameter element v[i][j]. I have tried several ways to do this using torch.autograd.grad() by backpropagating the first order gradient to each single element of the input velocity model, but it seems not working. My question is, is it possible to calculate gradient only one one single element of the input velocity model, e.g. something like gradient_00 = torch.autograd.grad(grad[0][0], v[0][0], create_graph=True)[0].

Another workaround maybe we can calculate an approximate illumination compensation, which requires wavefield at each time step. Is there is a way to access the wavefield at each time step for the forward modeling engine scalar?

Thank you very much!

[ar4]

Dear @shuhuah,

Thank you for your kind words. I am delighted that you find Deepwave useful.

Calculating the Hessian is unfortunately, as you say, very computationally expensive. Your idea about trying to approximate it using only its diagonal is good, but I think you are right that it probably won't be as easy as calling autograd.grad. In good news, you can indeed calculate an approximate illumination compensation, and I know that several people have done it using Deepwave.

If you only need the forward wavefields for the approach that you wish to use, then an easy way to do it is to put a receiver in each grid cell.

For better performance and lower memory requirements, method 2 in the custom imaging condition example shows a lower-level approach. The example modifies the backward propagator so that the forward and backward wavefields can be accessed in an imaging condition. If you only want to access the forward wavefields then you can instead modify the forward propagator. As in the backward propagator of the example, you should put a loop over time steps around the call to the C/CUDA forward function. This will allow you to access the wavefield after each time step with very low overhead.

Please let me know if you have further questions or need any more help with it.