How to verify if my added node features are propagating during train time

[varunpatro1]

Hello,

I’m working on a modification to the Allegro codebase that utilizes nequip 0.5.5. What I’d like to do is concatenate features along the dim=-1 of the one_hot vector and use this as part of the representation on each node. Based on reasoning from other posts, I did following in the _allegro.py file in allegro/model (just setting set_features to false, and providing irreps_in as in mir-group/nequip#360):

`def Allegro(config, initialize: bool, dataset: Optional[AtomicDataset] = None):
logging.debug("Building Allegro model...")

# Handle avg num neighbors auto
builder_utils.add_avg_num_neighbors(
    config=config, initialize=initialize, dataset=dataset
)

# Handle simple irreps
if "l_max" in config:
    l_max = int(config["l_max"])
    parity_setting = config["parity"]
    assert parity_setting in ("o3_full", "o3_restricted", "so3")
    irreps_edge_sh = repr(
        o3.Irreps.spherical_harmonics(
            l_max, p=(1 if parity_setting == "so3" else -1)
        )
    )
    nonscalars_include_parity = parity_setting == "o3_full"
    # check consistant
    assert config.get("irreps_edge_sh", irreps_edge_sh) == irreps_edge_sh
    assert (
        config.get("nonscalars_include_parity", nonscalars_include_parity)
        == nonscalars_include_parity
    )
    config["irreps_edge_sh"] = irreps_edge_sh
    config["nonscalars_include_parity"] = nonscalars_include_parity

layers = {
    # -- Encode --
    # Get various edge invariants
    "one_hot": (OneHotAtomEncoding, dict(set_features=False)),   # !!!
    "radial_basis": (
        RadialBasisEdgeEncoding,
        dict(
            basis=(
                NormalizedBasis
                if config.get("normalize_basis", True)
                else BesselBasis
            ),
            out_field=AtomicDataDict.EDGE_EMBEDDING_KEY,
        ),
    ),
    # Get edge nonscalars
    "spharm": SphericalHarmonicEdgeAttrs,
    # The core allegro model:
    "allegro": (
        Allegro_Module,
        dict(
            field=AtomicDataDict.EDGE_ATTRS_KEY,  # initial input is the edge SH
            edge_invariant_field=AtomicDataDict.EDGE_EMBEDDING_KEY,
            node_invariant_field=AtomicDataDict.NODE_ATTRS_KEY,
        ),
    ),
    "edge_eng": (
        ScalarMLP,
        dict(field=EDGE_FEATURES, out_field=EDGE_ENERGY, mlp_output_dimension=1),
    ),
    # Sum edgewise energies -> per-atom energies:
    "edge_eng_sum": EdgewiseEnergySum,
    # Sum system energy:
    "total_energy_sum": (
        AtomwiseReduce,
        dict(
            reduce="sum",
            field=AtomicDataDict.PER_ATOM_ENERGY_KEY,
            out_field=AtomicDataDict.TOTAL_ENERGY_KEY,
        ),
    ),
}

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

return model

I denoted the changed lines with !!! In my one_hot file in nequip, I have done the following:

concat = torch.cat((one_hot, atomic_numbers.to(type_numbers.device), enegs.to(type_numbers.device), radii.to(type_numbers.device), el_affs.to(type_numbers.device)), dim = -1)

To check whether my modifications are changing the final resuts, I also did a version of this where I concatenated 4 vectors with values on the magnitude of 1e7 (garbage values, essentially) to see if this would negatively impact the losses. I found that neither of my modifications (constructive or destructive) seem to significantly change the evaluation metrics with respect to the original Allegro on a particular dataset of conformers I have. (modified contains the added electronegativities, atomic numbers and radii, and garbage contains features of arbitrary large positive and negative numbers).

Is there a more direct way I can check to see whether my added features are being propagated? Or better yet, can I somehow retrieve the dependency of my losses on these added features? Any help would be great. Thank you.

[cw-tan]

Hi @varunpatro1,

Thank you for your interest in Allegro! Would it be possible to share how you've edited OneHotAtomEncoding (this module)? Also, perhaps you've done this already, but did you pip install again after your changes, or did you do pip install -e . so that your changes will be reflected without having to pip install each time after editing files in nequip? I assume you also did something like data[AtomicDataDict.NODE_ATTRS_KEY] = concat after you've done the concatenation? You could add print statements in nn/_allegro.py around here to check if the shape of node_invariants is what you expect after your concatenation.

[varunpatro1]

Hi @cw-tan,

Thank you for the quick reply! Here is my one_hot.py file, specifically the forward function for OneHotAtomEncoding. I have not altered the __init__.

`
def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:

    type_numbers = data[AtomicDataDict.ATOM_TYPE_KEY].squeeze(-1)
    one_hot = torch.nn.functional.one_hot(
        type_numbers, num_classes=self.num_types
    ).to(device=type_numbers.device, dtype=data[AtomicDataDict.POSITIONS_KEY].dtype)

    atomic_numbers = get_atomic_numbers(type_numbers=type_numbers)
    enegs = get_enegs(type_numbers=type_numbers)
    radii = get_atomic_radii(type_numbers=type_numbers)
    el_affs = get_el_affs(type_numbers=type_numbers)

    concat = torch.cat((one_hot, atomic_numbers.to(type_numbers.device), enegs.to(type_numbers.device), radii.to(type_numbers.device), el_affs.to(type_numbers.device)), dim = -1)

    data[AtomicDataDict.NODE_ATTRS_KEY] = one_hot
    if self.set_features:
        #print('set_features was true: features being set to original one_hot', flush=True)
        data[AtomicDataDict.NODE_FEATURES_KEY] = one_hot
    else:
        #print('set_features was false: features being set to concat', flush=True)
        data[AtomicDataDict.NODE_FEATURES_KEY] = concat

    return data`

To give an idea of the helper functions that generate the features to concatenate, here is the atomic numbers function below. The other functions are exactly alike save for the actual values to be concatenated.

`
def get_atomic_numbers(type_numbers):

# maps type_numbers to atomic_numbers
atom_mapping = {0: 1,
                1: 6,
                2: 9,
                3: 17,
                4: 35}

mapped_list = [atom_mapping[type_num.item()] for type_num in type_numbers]
atomic_numbers = torch.tensor(mapped_list).view(len(mapped_list),1)

atomic_numbers = atomic_numbers.float()
mean = atomic_numbers.mean()
std = atomic_numbers.std()

# Normalize the tensor
normalized_atomic_numbers = (atomic_numbers - mean) / std

return normalized_atomic_numbers`

Also, I did make sure to use the pip install -e . command so my changes would be transferred. I confirmed because I'm seeing the print statements I added to the one_hot.py file.

Additionally, regarding which dictionary field to add the concatenated features to, I wanted my added features to be static and just fed to the input OneHotAtomEncoding, I set them to data[AtomicDataDict.NODE_FEATURES_KEY] rather than data[AtomicDataDict.NODE_ATTRS_KEY]. Is this the correct way to add them? Adding them to data[AtomicDataDict.NODE_ATTRS_KEY] gives me a matrix multiplication error.

As for checking the shape in nn/_allegro.py, do you mean around line 500 where I should check the shape of new_latents? Thanks again for your help and I look forward to hearing from you!

[cw-tan]

Thanks for the information @varunpatro1, let me address a couple of points.

1. The one-hot encoding already encodes the atom types, so maybe having both one_hot and atomic_numbers may not be necessary? But that's a different discussion.
2. This is probably the main point - if I'm not mistaken, Allegro doesn't really use NODE_FEATURES_KEY, but it does use NODE_ATTRS_KEY (see here). So by setting concat to data[AtomicDataDict.NODE_FEATURES_KEY], it isn't going to be used anywhere in a meaningful way. The matrix multiplication error you got when you put concat in data[AtomicDataDict.NODE_ATTRS_KEY] is a sign that your changes have been propagated since the sizes of the "channel" or [-1] shape component of the node embedding tensor has changed. So, what you should do is set data[AtomicDataDict.NODE_ATTRS_KEY] = concat and make sure you register the dimensions correctly, i.e. look here. Currently, it says irreps_out = {AtomicDataDict.NODE_ATTRS_KEY: Irreps([(self.num_types, (0, 1))])}. You want to change self.num_types to whatever your concat.shape[-1] is, noting that one_hot.shape[-1] should be self.num_types. (this is an in-the-weeds detail, but that registration is important because it is used here in Allegro to determine the shapes of the multilayer perceptrons (MLPs) used, and that's why there was a matmul shape error - the weight matrix shape determined from here was likely just inconsistent with the shape of concat because it expected the shape of one_hot.)
3. For printing in nn/_allegro.py, I just meant here exactly, i.e. node_invariants = data[self.node_invariant_field] because node_invariants is probably the most reasonable (or maybe even the only way) to get the information you've put in concat into Allegro. And notice that during instantiation of the Allegro module, self.node_invariant_field would have been set be NODE_ATTRS_KEY by default (see here), i.e. node_invariant_field: str = AtomicDataDict.NODE_ATTRS_KEY in the __init__() args.

Hope this helps!

[varunpatro1]

Hi @cw-tan,

Thanks again for the prompt response.

1. Yes, this makes sense. Perhaps other atomic properties (electronegativity, atomic radius, mass etc) depending on the nature of the dataset might be more useful.
2. Ah, yes this would explain the similarity of all my results. I see how node_invariant_field: str = AtomicDataDict.NODE_ATTRS_KEY this line in allegro/nn/_.allegro.py means that my concat should be given to AtomicDataDict.NODE_ATTRS_KEY rather than AtomicDataDict.NODE_FEATURES_KEY. After making this modification, I did indeed get the same matrix multiplication error as before, and then what solved it was your comment about irreps: setting irreps_out = {AtomicDataDict.NODE_ATTRS_KEY: Irreps([(self.num_types+num_features_added, (0, 1))])}.
3. I printed out the shape where you suggested and saw that it reflected my concatenations as well.

Just to confirm, if I make these changes in one_hot.py, to both __init__ and forward, and the shape of my node_invariants is the concatenated shape, then I don't need to make any changes to the allegro/model/._allegro.py as I originally was based on the discussion in nequip? I'll rerun my old experiments and see if my values are hopefully changing now.

Thanks very much for all the help!

[cw-tan]

Just to confirm, if I make these changes in one_hot.py, to both __init__ and forward, and the shape of my node_invariants is the concatenated shape, then I don't need to make any changes to the allegro/model/._allegro.py as I originally was based on the discussion in nequip? I'll rerun my old experiments and see if my values are hopefully changing now.

Yes - the point of registering the correct shape in the __init__ in one_hot.py is so that later modules can access that shape information (in this case the Allegro module).

[varunpatro1]

Thanks a lot -- appreciate the help!

[Linux-cpp-lisp]

Just confirming from the above the NODE_FEATURES_KEY is not used by Allegro. Will be curious to hear @varunpatro1 if you find adding these atomic properties to be helpful in learning your data.

[varunpatro1]

Thank you for the confirmation @Linux-cpp-lisp. I'm getting very slight improvements in my early testing phases, and will be happy to share results if performance gets a bit better!

[Linux-cpp-lisp]

Great, sounds good!

[varunpatro1]

Hi, upon some more training with larger subsets of QM9, it seems like the additional concatenated features aren't in fact improving performance. Just wanted to let the authors/maintainers know. Best of luck and thanks for the discussions!

[cw-tan]

Cool, thanks for the update!