BatchMolGraph padding makes inference embeddings depend on batch composition

Summary

The representation of a molecule changes materially depending on which other
molecules are present in the same BatchMolGraph, even with model.eval() and
bond_drop_rate = 0.

For the first molecule in the panel below, comparing a singleton forward pass
with a forward pass containing all eight molecules gives a 7.402% relative L2
difference in the concatenated molecular fingerprint. Individual atom/bond
representations differ by roughly 2–11%.

This is deterministic: repeated singleton calls match each other, and repeated
eight-molecule calls match each other.

Tested against main at
8828743036675d1b6d5f4586ef7bcfea70233d59 with the released
kermt_contrastive_v2.0.pt checkpoint.

Expected behavior

A molecule's inference representation should be independent of unrelated
molecules collated into the same batch, apart from normal floating-point
accumulation noise.

Reproduction panel

SMILES = [
    "O=C1NC(=O)C(=Cc2ccc(-c3ccc(I)cc3)o2)C(=O)N1",
    "O=C(NCCOC(=O)[C@@H]1CCCCN1S(=O)(=O)Cc1ccccc1)c1cccnc1",
    "CCOC(=O)c1[nH]c(C)c(CCC(=O)N2CCN(c3ccc(OC)cc3)CC2)c1C",
    "CC(C)C[C@H](C[C@H](N)C(=O)O)C(=O)O",
    "CCCOc1ccc(S(=O)(=O)N2CCN(C)CC2)cc1-c1nc2c(ncn2CCC)c(=O)[nH]1",
    "CCOc1ncccc1C(=O)Nc1ccc(C(N)=O)cc1",
    "CC(C)(C)C(=O)OCc1ccc(C2=CC(=O)CC(C)(C)C2=O)cc1",
    "CCCOc1cc(OCCC)cc(-c2ccc(C(=O)Nc3ccc4ccc(C)nc4c3)c(=O)[nH]2)c1",
]

The model setup and forward calls were:

import copy
import torch

from kermt.data.molgraph import mol2graph
from kermt.model.models import KERMTEmbedding

checkpoint = torch.load(
    "kermt_contrastive_v2.0.pt",
    map_location="cpu",
    weights_only=False,
)
args = copy.copy(checkpoint["args"])
args.use_cuikmolmaker_featurization = False
args.bond_drop_rate = 0
args.no_cache = True
args.cuda = False

model = KERMTEmbedding(args)
model.load_state_dict(
    {
        key.removeprefix("kermt."): value
        for key, value in checkpoint["state_dict"].items()
        if key.startswith("kermt.")
    },
    strict=True,
)
model.eval()

def forward(smiles):
    graph = mol2graph(smiles, shared_dict={}, args=args)
    components = graph.get_components()
    with torch.no_grad():
        output = model(components)
    return components, output

singleton_components, singleton = forward(SMILES[:1])
batch_components, batch = forward(SMILES)

Slice molecule 0 using components[5][0] for atom rows and
components[6][0] for directed-bond rows. The relative error used here is:

torch.linalg.vector_norm(singleton_value - batch_value) / \
    torch.linalg.vector_norm(batch_value)

Observed relative L2 differences:

Representation	Relative L2
atom_from_atom	6.552%
atom_from_bond	11.172%
final node output	2.208%
bond_from_atom	6.053%
bond_from_bond	11.142%
concatenated molecular fingerprint (mean-pooled four branches)	7.402%

Suspected cause

BatchMolGraph determines max_num_bonds over the entire batch and pads all
neighbor lists to that width with index 0:

• https://github.com/NVIDIA-BioNeMo/KERMT/blob/8828743036675d1b6d5f4586ef7bcfea70233d59/kermt/data/molgraph.py#L364-L395

In this panel, the first molecule's singleton maximum degree is 3, while
another molecule contains a degree-4 sulfone sulfur. Adding it therefore
appends another index-0 padding entry to neighbor rows in the first molecule.

The sentinel feature starts as zero, but the encoder applies biased linear
projections and affine LayerNorm to the complete feature tensor, including row
0:

• https://github.com/NVIDIA-BioNeMo/KERMT/blob/8828743036675d1b6d5f4586ef7bcfea70233d59/kermt/model/layers.py#L636-L647

The transformed sentinel can consequently be nonzero.
select_neighbor_and_aggregate gathers the index-0 entries and sums them
without a padding mask:

• https://github.com/NVIDIA-BioNeMo/KERMT/blob/8828743036675d1b6d5f4586ef7bcfea70233d59/kermt/util/nn_utils.py#L101-L111

As a result, changing the batch-wide padded width changes every affected
aggregation.

Possible remediation

Mask index-0 padding before neighbor reduction, or otherwise keep each
molecule's effective neighbor width/sentinel behavior independent of its batch
companions.

A straightforward index-0 mask changes the historical singleton outputs too,
because existing checkpoints were trained with the unmasked sentinel behavior.
If checkpoint compatibility is required, retaining each molecule's singleton
effective width inside a packed batch avoids the batch-composition dependency
while preserving current singleton results.

I can provide a complete executable comparison script or downstream regression
test if helpful.

Start with the neighbor padding in kermt/data/molgraph.py, then trace the encoder and reduction paths in kermt/model/layers.py and kermt/util/nn_utils.py. Run the supplied eight-molecule reproduction and add regression coverage showing that a molecule's inference representation is unchanged by batch companions, while checking the stated checkpoint-compatibility concern.