Implementing dropout during training

[davidleocadio]

Hello,

It seems like NequIP doesn't currently support dropout. What is the right way to implement this in the source code? Also, which layers would dropout be applicable to?

From a previous discussion it seems like dropout should happen in:
- certainly in the invariant radial network MLPs
- certainly not in the chemical embedding
- certainly not in any TPs
- to discuss: in the equivariant linear self-interaction layers (shouldn’t mess with equivariance)

Thanks!
David

[Linux-cpp-lisp]

Hi David,

The radial networks in NequIP are implemented as e3nn.nn.FullyConnectedNet. If you want dropout once before the entire radial network, you can add that in InteractionBlock. If you want dropout at every layer of the radial network, you'd need to make a PR to e3nn adding that feature as an option to FullyConnectedNet (you'd basically add a torch.nn.Dropout here: https://github.com/e3nn/e3nn/blob/main/e3nn/nn/_fc.py#L84-L85).

The other option as you mention is dropout in the equivariant feature vectors. Equivariant dropout is trickier: you can't just dropout elements from the tensor without breaking symmetry. Instead, you have to drop out/retain entire irreps. There is an implementation of this kind of dropout in e3nn (https://github.com/e3nn/e3nn/blob/main/e3nn/nn/_dropout.py) but I've never used it.

You could put equivariant dropout in a lot of places. The most obvious place to me would be between convolution layers, but it could go anywhere where equivariant data is propagated.

To add something like this to NequIP, I would take advantage of the nequip framework's modular architecture to implement a Dropout module wrapping e3nn's. An example of what a simple NequIP module looks like is https://github.com/mir-group/nequip/blob/main/nequip/nn/_atomwise.py#L14. You'd then add the Dropout module to the definition of the model between the ConvNetLayer modules (https://github.com/mir-group/nequip/blob/main/nequip/models/_eng.py#L41).

I would accept a PR implementing either/both of these features if you are interested! 😃

[Linux-cpp-lisp]

@davidleocadio yep, that's the code I was thinking of — what you'll need to do is wrap it in a NequIP Dropout module that acts on AtomicDataDict.Type, rather than on a single tensor directly. (See the AtomwiseLinear example I linked above for what wrapping a simple e3nn operation looks like.) This is much better for modularity than adding one more option to, say, ConvNetLayer.

We have a small system that automatically matches options from the YAML file to __init__ parameters of modules that are built as part of the model. So all you'd need to do is to define a parameter, say radial_dropout_rate in __init__ of InteractionBlock (which builds the e3nn.nn.FullyConnectedNet to use internally). This parameter is then automatically populated by YAML keys radial_dropout_rate and InteractionBlock_radial_dropout_rate (in increasing order of preference).

Similarly, for adding a Dropout module, you'd also just add dropout_rate (or whatever) as a __init__ parameter.

[Linux-cpp-lisp]

Also, in no particular order:

• Please branch off of develop, or just be aware that you'd need to rebase to it before merging a PR
• Don't add Dropout modules to the main _eng.py, since that is the standard official model and we try not to change that much. Instead, copy the model definition into your own file, add the dropouts, and maybe add this file to the examples/ folder
• While Dropout is a useful module to have in the repo no matter what, I'd prefer to have some sign that radial net dropout does something useful (or at least not counterproductive) before adding it to InteractionBlock, since that's a part of every model no matter what.

[Linux-cpp-lisp]

Eventually all this information will make it into CONTRIBUTING.yaml... for now, it's all here 😄

The developer tutorial notebook may also be of some use: https://deepnote.com/project/NequIP-Tutorial-JBLKk3rRRFiXLF1a3Vpmfg/%2Ftutorial_.ipynb/#00034-10c1dfd2-e4e3-4b66-b892-61fa5cee872e

[Linux-cpp-lisp]

I would strongly encourage you to experiment quickly locally with radial network dropout before investing any of your time into making PRs for it, since given the extremely small size of our typical radial networks, along with everything I said before, I don't think it's a given that dropout there will be useful

[Linux-cpp-lisp]

And finally, of course, thanks for your interest in our project, and in contributing back to the common code with your work!!