fix(model): zero masked pair slots before the reverse gather

[E-Rum]

Problem

Backbone.__call__ masks edge features going into the transformer layers, but
TransformerLayer returns edge + attn + mlp(...) unmasked, so padded pair slots
leave each GNN layer holding nonzero features. The reverse gather
(reversed_flat = edge_flat[reverse]) is the one place a masked slot is read, so
any reverse entry pointing at the padded sentinel picks those features up — and
the model's output then depends on the packed width k, which is a caller-side
shape choice, not physics.

Within pet-jax's own selection path this is unreachable in normal operation:
_select_edges is symmetric (pair_cutoffs averages the two atoms' cutoffs), so
a selected pair's reciprocal is always selected too, and a sentinel-pointing
reverse arises only under overflow — already a flagged failure. But
pack_edges accepts a pre-trimmed neighbour list and documents only a per-center
capacity requirement, so a caller that trims per-center by distance rank (an
asymmetric cut, unlike a distance cut) hits this in steady state with
overflow=False and no diagnostic.

Measured on a downstream trained checkpoint whose neighbour list is knn-trimmed
per center, where ~10% of surviving directed edges have lost their reciprocal:

pack width	valid E RMSE	valid F RMSE
max_neighbors + 1 (its training width)	baseline	baseline
192	9.2x worse	3.3x worse

overflow is False in both. Force error degrades monotonically across widths
between those two points. Bisecting the batch axes one at a time, every other
padded axis is exactly invariant — only the pack width moves the result. With a
symmetric neighbour list (reciprocal-pruned, no orphans) the two widths agree to
4e-6 eV, which isolates the orphan gather as the sole path.

The width dependence comes from where the sentinel lives. It is the last slot of
the padded dummy-atom row, so that row is fully masked; its center token is
Dense(0) = bias, nonzero, while its edge tokens are zero, so the attention
average puts weight 1/(1 + k) on a nonzero vector. On the checkpoint above the
sentinel feature vector drifts in norm from 8.75 at k=17 to 9.77 at k=384,
saturating, and every orphan edge gathers it.

Fix

-            edge_flat = edge.reshape(P, d_pet)
+            edge_flat = edge.reshape(P, d_pet) * pair_mask[..., None]

One multiply per GNN layer. Masked slots now leave the backbone at exactly zero,
which makes the forward invariant to the pack width, makes the sentinel the
zeroed edge that callers of pack_edges reasonably assume, and also cleans up
pet-jax's own overflow path.

Tests

tests/test_select.py:

• test_masked_slots_carry_no_features — padded slots leave the backbone at
  exactly zero. This is the guard: it fails on main at plain init, with no
  tolerances involved.
• test_backbone_invariant_to_packed_width — with an orphaned reverse, features
  at k and k + 11 must agree.
• _orphaned() helper — drops one direction of one pair, reproducing an upstream
  per-center trim.

Both fixture cases (rocksalt crystal, H2O): 4 failures on main, 0 with the fix.
The existing test_pack_edges_round_trip / test_pack_edges_reverse_involution
pass throughout — note the latter asserts the very symmetry this bug hides behind,
and passes only because to_structure yields a symmetric neighbour list.

Two things worth knowing if you extend these tests:

• A fresh init leaves every bias zero, which zeroes the sentinel and hides the
  bug completely. The invariance test shifts parameters off init deliberately.
• An energy-level invariance check on a small randomly-initialised model does
  not discriminate — pre/post agreement stayed at ~1e-9 relative even with bias
  offsets up to +8. The leak only becomes load-bearing in a trained model, where
  the junk vector is comparable in magnitude to real edge features. Assert the
  structural property, not downstream energies.

Compatibility

This changes the function for any checkpoint trained through an asymmetric
neighbour list: those weights have absorbed the sentinel as a de facto constant
input. Applying the fix to such a checkpoint without retraining degraded its
validation error by roughly two orders of magnitude, so affected checkpoints need
retraining rather than a reload.

Checkpoints with no orphan edges are unaffected: a run whose trim never binds
(max_neighbors below the cap, so the neighbour list stays symmetric) is
bit-identical before and after.

Not addressed here

• pack_edges still accepts an asymmetric neighbour list silently. Its
  overflow flag reports only per-center capacity, not "a selected pair's
  reciprocal was dropped". Worth either documenting the symmetry requirement or
  extending the flag, so the next caller doesn't repeat this.
• Unverified adjacent hunch: in _pack_selected_to_flat, every non-fitting pair
  writes to slot P_sel - 1 via sel_to_pair.at[slot].set(...), which would
  collide with a real edge if the last atom row were ever exactly full. Trailing
  padded rows appear to prevent it; I did not construct a case.
• CI should confirm the reference-prediction and calculator suites, which I could
  not run locally (git-lfs assets, and no pytest in my container — I drove the
  test functions directly).

🤖 Generated with Claude Code