Add hooks for indirect learning. - #1209

Seems to work :)

Thanks!

One question though: is it a hook job to take per atom prediction and make it per-system? It kinda look like this is the case here? Or I might have misunderstood the example

There could be a hook to do that, specially for cases when the aggregation is not simply a sum, but something more complicated like softmin/max for the bandgaps. I guess the case of the global dipoles can be seen as that, we are doing dipole = sum(local_charge * positions + local_dipole), so it's not just a simple sum.

I guess indeed with standardization of outputs these hooks might become applied by default when needed, instead of the user having to specify them.

@cesaremalosso I have added the hook to produce gap-like targets. I have been testing it and seems to work with both PET (thanks to Joe) and MACE, give it a test when you can :) I basically copied the functionality that you had in your branch, and now there's a hook called minmax_gap. Don't worry the name is not definitive haha

The way of using it is:

architecture:
  name: pet | experimental.mace
  model:
    ...other model hypers
    post_hooks:
      minmax_gap: mtt::gap_energy
  training:
    atomic_baseline: {mtt::gap_energy: 0.0}
    ...whatever

training_set:
  systems: train_structures.xyz
  targets:
    mtt::gap_energy:
      quantity: energy
      key: Gap_energy
      sample_kind: system
      type: scalar
      forces:
        key: Gap_forces

also you can tune the hook further with something like:

post_hooks:
  minmax_gap:
    outputs: mtt::gap_energy
    pooling: {type: softmax}

Let me know if it works fine!

Hmm for the gap case the composition model needs to be disabled, I wonder if this should be the job of the hook to warn about this. It certainly complicates things 😅

Also Cesare is asking if one could do finetuning of an MLIP by adding these hooks, and I see no easy way, since right now the hooks are specified as model hypers, which are not used when finetuning I think (?)

So maybe the right design of a UI would be that users specify these hooks in the target itself:

targets:
  mtt::gap_energy:
    ...target specification
    quantity: intensive_gap # Or whatever way of specifying it.

this would make it cleaner for the composition model to avoid this target and for the finetuning to detect it, and probably also be in the spirit of metatrain of "we should require as litle knowledge from users".

But then some cases become more complicated, like when a hook has hyperparameters or when a hook produces multiple outputs (e.g. a diagonalization hook that produces eigenvalues, eigenvectors, density matrices, and other electronic structure properties). In general, passing the category in the targets results in a less flexible UI I think, for the good and for the bad 😅

Geeez, this looks very complicated, but also very useful. So the idea is that you can specify generic modules to be applied after a head? I think that defining them in the targets as post_hooks could be a good solution no? or does that interfere with mtt understanding what the actual outputs are?

In essence it is what you are saying, and very similar to the approach of metatomic to define wrappers, except these hooks work also during training. Now the problem is that I don't know what is the best way to expose it to the users.

In the way metatrain does things, I believe it has to be fully controlled by the architecture. So specifying a full hook in the target seems a bit strange, as some architectures might not want to use that hook (for example a graph2mat hook for edges, PET will not want to use it). I think it is more metatrain-like to do:

• Specify quantity in the target.
• Architecture decides that for this quantity it is going to offload the work to hooks.

But then you lose quite some flexibility (do we want that flexibility?). Specifying the full hook in the target would indeed require modifying metatrain.TargetInfo and metatomic.ModelOutput, so it seems it would be a more involved path. But maybe it is the way to go, I don't know. This is why I'm trying to gather examples that would use the functionality, to see if there's a simple way to make them all work.

Ok so with this we have the three examples that I had in mind for hooks:

• Softmin/max for gaps (custom pooling from atomic contributions).
• Global multipoles (pooling from atomic contributions involving the systems' positions).
• Tensor basis (hook with learnable parameters).

The other two that I had involve matrices (graph2mat hook and eigenvalues hook) so they can't be added yet. If someone has more use cases let me know and let's add them. Then once we have all the examples I'll start refining the implementation.

I think the matrix use case is sufficiently complex and useful to consider here in the design (even if not directly implementable yet from the current state of main) as it is different from the others: for the Hamiltonian, for example, we might want to directly and indirectly learn on it simultaneously (i.e. H in the loss function as well as derived eigs and DM), and the basis set still needs passing somehow to the model. This is different from the total dipole, for instance, where the dimensions of the indirectly learned charges and vector contributions are predetermined, there is no reference data for them, and they only enter the loss as a derived quantity (i.e. no term explicitly on them)

Yes, the basis information will be passed through dataset_info, and it will be possible to use targets as inputs for the hooks

the basis information will be passed through dataset_info

Ok good good.

A small side note to be aware of: we need to make sure that the hooks are compatible with the optional dependencies installed by the user. Currently when I install with just the base dependencies (i.e. for PET) I get a spex import error from the tensor basis hook module.

Yeah I thought about this, probably optional dependencies for hooks should be in pyproject.toml as [hook-*]

Ok so this got x10 more complex to review haha it's probably best to explain in person.

Essentially we realised that at some point we would need checkpoint updates, things need to be properly documented, well tested, etc... so I decided to follow the same design as with architectures (folder for each hook, automatic documentation, common tests...).

Still not finished, but decided to push so that people can modify the hooks in this new setup.

I'm happy because the hooks seem like a very nice way to release code when people publish papers, which has been a bit of a problem lately since modifying an architecture is not easy. The barrier for making it serious enough for a hook is lower and there is no friction with architecture maintainers.