Analytical Hessian in ase calculator about mace HOT 9 CLOSED

That should be pretty trivial, just add to utils.py a loop of grad. Might need to do some matrix reconstruction though.

from mace.

I don't think analytical Hessian belongs in the ASE calculator, at least for now - it's useful, but it's not a standard ASE calculator result. It is trivial to do naively, or you can use https://wiki.fysik.dtu.dk/ase/ase/vibrations/modes.html, or use phonopy to include symmetry to reduce the number of finite differences.

from mace.

I am aware that there is not too much of an interested in implementing Hessians, but a proper implementation would be very useful. Often one of the selling points of MLIP are the easy to get Hessians and therefore there should be a community interested in that. So was trying to implement hessians for MACE and I was hitting some problems. As recommended from @ilyes319 a loop is a way to implement hessian:

hessian1=[]
for grad_elem in forces.view(-1):
            print(grad_elem)
            hess_row = torch.autograd.grad(
                outputs = [grad_elem], 
                inputs = [positions], 
                grad_outputs=torch.ones_like(grad_elem),
                retain_graph=True, 
                create_graph=False,
                allow_unused=False,
            )[0]
            hess_row =hess_row.detach() #this makes it very slow? but needs less memory
            if hess_row is None:
                hessian1.append(torch.zeros_like(positions))
            else:
                hessian1.append(hess_row.view(1, -1))

This works very good!, but problem here is, for larger systems it takes very long to iterate through, here time is the limiting factor.

It is also possible to make it vectorized, which is significantly faster but also needs a bit better understanding of the autograd.grad function:

tuple_of_tensors=torch.eye(forces.view(-1).shape[0], forces.view(-1).shape[0],device="cuda")
hessian= torch.autograd.grad(
            outputs = [forces.view(-1)],#*45,# forces.view(-1), 
            inputs = (positions,),
            grad_outputs=(tuple_of_tensors,),
            retain_graph=True,
            create_graph=False,
            allow_unused=False,
            is_grads_batched=True,
        )

Here the is_grads_batched allows a vectorized behavior, but this has the problem of being extremely memory intensive.
I also tried implementations based torch.autograd.functional.hessian ortorch.autograd.functional.jacobianor the torch.func.hessian function with rewritting larger parts of MACE into a function that yields energy and forces respectively.
Nevertheless, non of these implementations is memory efficient and fast at the same time
For systems with more than 100 atoms they are either to slow or to memory intense. Maybe someone here, with a better understanding of the MACE code and pytorch knows a solution for fast and memory efficient Hessians for MACE.

from mace.

Indeed computing analytical Hessian is quite expensive. I think these are mostly all the exact methods from torch. Can you give me an estimate of the sizes of systems you want to look at?
Also can you tell me the GPU you are using and the timing for each of the methods.

from mace.

I am using a A100 with 80Gb RAM
My current testing code looks like this:

from mace.calculators import mace_mp
from ase import build

atoms = build.molecule('H2O')
calc = mace_mp(model="medium", dispersion=False, default_dtype="float32",device='cuda' )
atoms.calc = calc
#print(atoms.get_potential_energy())

#exit()

from ase.build import fcc111

initial = fcc111('Pt', size=(7, 8, 1), vacuum=10.0, orthogonal=True) # size i would like to test = (7,8,6)
initial.calc = calc

e0=initial.get_potential_energy()
print(e0)

For the loop approach this takes ~ 4.5 seconds (could restructure my code then I could give you a more precise answer).
For the vectorized approach this takes ~1.5 seconds.
With the size of the system I am actually working with (for example size = (7,8,6))
loop: ~85 seconds
vectorized: Memory overflow
this is ONLY the hessian calculation by the autograd,grad!

If i use autograd.functional. hessian:

h = torch.autograd.functional.hessian(
                self.get_energy,
                inputs=data["positions"],
                vectorize = False,
                #outer_jacobian_strategy= 'reverse-mode'
            
            )

System size = (7,8,1), vecorize =False: ~6.8 seconds
System size = (7,8,1), vecorize =True: ~1.8 seconds

System size = (7,8,6), vecorize =False: ~87 seconds
System size = (7,8,6), vecorize =True: Memory overflow

h = torch.autograd.functional.jacobian(self.get_forces,data["positions"], vectorize=False)
System size = (7,8,1), vecorize =False: ~6.8 seconds
System size = (7,8,1), vecorize =True: ~1.8 seconds

The two torch.func approaches

h=torch.func.hessian(self.get_energy)(data["positions"])#,data["node_attrs"])
h=torch.func.jacrev(torch.func.jacfwd(self.get_energy),chunk_size= 1)(data["positions"])#,data["node_attrs"]) 

are both very similar in speed ~1.3 seconds but sometimes have already a memory overflow for the system size (7,8,1).

I am overall not the biggest fan of rewritting too much of the code so all options with get_energy or get_forces are rather experimental.

Thanks for the fast answer I hope this stuff helps.

from mace.

One thing I haven't noticed in this discussion is whether you're trying to compute the entire Hessian matrix explicitly, or whether you're using symmetry to take advantage of the fact that all atoms are equivalent in your fcc cell and you only need $dF_i / dr_1$, with only $i$ looping over all the atoms. Of course in a disordered system this won't be the case, but for a lot of solid state situations that'd be sufficient, and much cheaper (albeit harder to assemble into the complete matrix - that's something phonopy does automatically, e.g.)

from mace.

While my testing system is a fcc cell, I will also need Hessian matrices of disordered systems and therefore I need them explicitly. Ofc an estimation method could be sufficient but the reason why I want to use a MLIP is this idea of cheap and fast automated derivatives.

from mace.

Have you tried batching the rows, so that it does a few at a time, minimizing overhead, but doesn't exceed memory?

from mace.

I indeed tried to do it in batches and it seems to be the best solution for now but it doesn't feel to like a good solution but maybe it is the best for now.

Edit: And it is also still very inefficent

from mace.