Dealing with additional node features and derivatives w.r.t. them

[shasax]

We are specifically dealing with atomic clusters of around 40-50 atoms (single species for now) where the atoms have something in addition to their positions as an input feature. The output we need to learn is the cluster energy and force on each atom. Here are a couple of questions I have in this regard:

1. To provide 'position variances' as an additional input feature on each node, is it okay to use the keyword 'node_features' in 'key_mapping' in the config file?

2. We need to train our model to obtain the forces, which are the energy derivatives w.r.t. positions, and also to obtain the derivatives w.r.t. the additional input node features we have. Is there any way/field in the config file to make this possible?

[Linux-cpp-lisp]

Hi @shasax ,

Sounds like an interesting project!

Question 1: Almost. Providing that key_mapping will give it as input to the model, but you also need to:

1. Decide whether you want to concatenate your extra input to the existing one-hot input for node_features or replace it
2. Clarify what the irreps for your "position variances" are, and provide them appropriately
3. Tell the OneHotAtomEncoding module in the model not to override your provided node_features

The correct way to do this is to define a new "model builder": a function that builds/updates the model in nequip. By default, nequip uses the following list of model builders, which I'm showing here as you would specify them in the YAML config file:

model_builders:
  - SimpleIrrepsConfig
  - EnergyModel
  - PerSpeciesRescale
  - StressForceOutput
  - RescaleEnergyEtc

You will want to replace EnergyModel with a model builder that creates a model with the extra input feature. You do this by defining a simple Python function in a place where it can be imported (either a setuptools installed package or some hacks with PYTHONPATH like in nequip/examples/lj). Here's an appropriately updated version of the default model builder in https://github.com/mir-group/nequip/blob/main/nequip/model/_eng.py:

def MyFancyEnergyModel(
    config, initialize: bool, dataset: Optional[AtomicDataset] = None
) -> SequentialGraphNetwork:
    logging.debug("Start building the network model")

    builder_utils.add_avg_num_neighbors(
        config=config, initialize=initialize, dataset=dataset
    )

    num_layers = config.get("num_layers", 3)

    layers = {
        # -- Encode --
        "one_hot": (OneHotAtomEncoding, dict(set_features=False)),   # !!!
        "spharm_edges": SphericalHarmonicEdgeAttrs,
        "radial_basis": RadialBasisEdgeEncoding,
        # -- Embed features --
        "chemical_embedding": AtomwiseLinear,
    }

    # add convnet layers
    # insertion preserves order
    for layer_i in range(num_layers):
        layers[f"layer{layer_i}_convnet"] = ConvNetLayer

    # .update also maintains insertion order
    layers.update(
        {
            # TODO: the next linear throws out all L > 0, don't create them in the last layer of convnet
            # -- output block --
            "conv_to_output_hidden": AtomwiseLinear,
            "output_hidden_to_scalar": (
                AtomwiseLinear,
                dict(irreps_out="1x0e", out_field=AtomicDataDict.PER_ATOM_ENERGY_KEY),
            ),
        }
    )

    layers["total_energy_sum"] = (
        AtomwiseReduce,
        dict(
            reduce="sum",
            field=AtomicDataDict.PER_ATOM_ENERGY_KEY,
            out_field=AtomicDataDict.TOTAL_ENERGY_KEY,
        ),
    )

    return SequentialGraphNetwork.from_parameters(
        shared_params=config,
        layers=layers,
        irreps_in={"node_features": "4x0e"} #!!!
    )

I've marked edited lines with #!!! comments. This is the version for if you want to replace the input one-hot encoding with your input feature. There are two edits:

1. Telling OneHotAtomEncoding not to override the features
2. Specifying the expected irreps for the input features (at the very end, you'd obviously have to change the irreps to be correct)

If you want to do the version of this where you concatenate the one-hot with your features, you'd add another module to the layers in the builder above to do so.

Question 2: This is not possible from the config file alone, but is possible in a straightforward and supported way again using model builders. With the same set up as before, we can modify the existing model builder for forces (https://github.com/mir-group/nequip/blob/main/nequip/model/_grads.py) to wrap the model:

def NodeFeaturesGradOutput(model: GraphModuleMixin) -> GradientOutput:
    return GradientOutput(
        func=model,
        of=AtomicDataDict.TOTAL_ENERGY_KEY,
        wrt=AtomicDataDict.NODE_FEATURES_KEY,
        out_field="node_feature_forces",
        sign=-1,  # force is the negative gradient
    )

Then you just add this builder to your list in your YAML config:

model_builders:
  - SimpleIrrepsConfig
  - mystuff.MyFancyEnergyModel
  - PerSpeciesRescale
  - ForceOutput
  - mystuff.NodeFeaturesGradOutput
  - RescaleEnergyEtc

[Linux-cpp-lisp]

Can you elaborate on what numerical differentiation you are comparing to?

Something might be wrong with the scaling of the node feature forces... I may have given you the wrong advice there. Since they are a derivative of the energy, by linearity, they also have to share that same scaling--- otherwise it doesn't make any sense. So really, NODE_FEATURE_FORCES should be added to ALL_ENERGY_KEYS so that it gets rescaled in tandem with the energy, forces, stress, etc. by the normal force RMS.

That still doesn't quite resolve the issue of the scale of the input node features, but that's something that has to be solved in preprocessing of the dataset or by scaling them to normal by some constant as a first step in the model inference pipeline. (Which is the opposite of scaling the output node feature forces, as is happening here.)

[shasax]

So, by numerical derivative, I mean that let's say we predict the model outputs on an input of same positions but uniformly spaced node_features.
In that case, at any input point, the numerical derivative of the predicted energy landscape wrt node_features (delta E / delta node_features) should be similar to the node_feature_forces we observe in the predictions. But these two are very much apart, which points to something fishy.

Ah, if that is the case for node_feature_forces scaling, I will try that. Previously, I was trying to scale them separately with their RMS value.

Regarding scale of input node features, I do them in the data preprocessing step now.

Please let me know any other implementations you might suggest to be wrong currently. Thanks again for the help :)

[Linux-cpp-lisp]

Got it, sounds correct--- I think it is the difference in the output scaling being inconsistent with the derivative then. Since d (scaling * E) / d node_features = scaling * dE/d node_features, but not equal to scaling2 * dE/d node_features.

Regarding scale of input node features, I do them in the data preprocessing step now.

Great, then you should be fine. (Be careful of course to scale the forces back out in post-processing when you do inference then.)

Please let me know any other implementations you might suggest to be wrong currently. Thanks again for the help :)

All sounds good to me right now 👍

[shasax]

Hello,
Thank you for all the support.
I am reopening this thread after a long time. So just to recap a bit, I have been working with NEQUIP after tweaking the EnergyModel slightly, where I replaced the one-hot atom encoding with a list of double-type node features. I previously received help from your side, which made me define my own energy model as shown below:

def MyFancyEnergyModel(
    config, initialize: bool, dataset: Optional[AtomicDataset] = None
) -> SequentialGraphNetwork:
    logging.debug("Start building the network model")

    builder_utils.add_avg_num_neighbors(
        config=config, initialize=initialize, dataset=dataset
    )

    num_layers = config.get("num_layers", 3)

    layers = {
        # -- Encode --
        "one_hot": (OneHotAtomEncoding, dict(set_features=False)),   # !!!
        "spharm_edges": SphericalHarmonicEdgeAttrs,
        "radial_basis": RadialBasisEdgeEncoding,
        # -- Embed features --
        "chemical_embedding": AtomwiseLinear,
    }

    # add convnet layers
    # insertion preserves order
    for layer_i in range(num_layers):
        layers[f"layer{layer_i}_convnet"] = ConvNetLayer

    # .update also maintains insertion order
    layers.update(
        {
            # TODO: the next linear throws out all L > 0, don't create them in the last layer of convnet
            # -- output block --
            "conv_to_output_hidden": AtomwiseLinear,
            "output_hidden_to_scalar": (
                AtomwiseLinear,
                dict(irreps_out="1x0e", out_field=AtomicDataDict.PER_ATOM_ENERGY_KEY),
            ),
        }
    )

    layers["total_energy_sum"] = (
        AtomwiseReduce,
        dict(
            reduce="sum",
            field=AtomicDataDict.PER_ATOM_ENERGY_KEY,
            out_field=AtomicDataDict.TOTAL_ENERGY_KEY,
        ),
    )

    return SequentialGraphNetwork.from_parameters(
        shared_params=config,
        layers=layers,
        irreps_in={"node_features": "4x0e"} #!!!
    )

However, now I would want to concatenate the one-hot atom encoding with my double-style node features. Can it be done using the same model builder? Or does this require significant changes?

What I believe would work is that I concatenated the one-hot tensor to my input node features inside the forward definition of class OneHotAtomEncoding, by using:

total_node_features = torch.cat((data[AtomicDataDict.NODE_FEATURES_KEY],one_hot),1)
data[AtomicDataDict.NODE_FEATURES_KEY] = total_node_features

Is something else also needed?

Also, I have a more extended version of the ForceOutput model builder, which computes the gradient w.r.t. the positions and the node features. This is shown below.

def CombinedGradOutput(model: GraphModuleMixin) -> GradientOutput:
    if AtomicDataDict.FORCE_KEY in model.irreps_out:
        raise ValueError("This model already has force outputs.")
    return GradientOutput(
        func=model,
        of=AtomicDataDict.TOTAL_ENERGY_KEY,
        wrt=[AtomicDataDict.NODE_FEATURES_KEY,AtomicDataDict.POSITIONS_KEY],
        out_field=[AtomicDataDict.NODE_FEATURE_FORCE_KEY,AtomicDataDict.FORCE_KEY],
        sign=-1,  # force is the negative gradient
    )

As the gradient-output builder is usually on the outside of the one-hot encoding, the node features are not yet concatenated with the one-hot features when the gradient is computed. But this also is something desired as I do not want derivatives w.r.t. the one-hot features.

Does this sound alright, or something needs to be changed also in the CombinedGradOutput?

Some help would be highly appreciated!

[Linux-cpp-lisp]

Hi @shasax ,

I think you're right that a modification / replacement of OneHotAtomEncoding to allow concatenating to existing features. Note, however, that you will also need to report the correct concatenation of the irreps of the one hot features and the incoming other node features in irreps_out of the OneHotAtomEncoding. (If I remember right, o3.Irreps, being Tuples, can be concatenated with +.) (Also can pay to be careful with dimensions, usually better to use -1 as the concat dim rather than 1 just to be on the safe side and consistent with other code in the framework.)

Regarding gradients, the GradientOutput happens around the entire SequentialGraphNetwork, so PyTorch can understand that the concatenated features are derived from the original node features and should propagate gradients accordingly. (In other words, you are asking for the gradient w.r.t. the node features when they enter the beginning of the GradientOutput, not at any later redefinition.)