PET attention: keys with cutoff factor exactly 0 are penalised, not excluded

[sirmarcel]

Component: PET (src/metatrain/pet), shared by FlashMD / Symplectic FlashMD
Revision inspected: main@2618bb0
Severity: exact-invariant violation; measured end-to-end effect on a trained checkpoint is negligible (see Measured impact). Filed for correctness/reproducibility, not as an accuracy fix.

Marcel note: This issue is a result of @MichelangeloDomina's recent bug hunting spree. Filed here by trusty Claude Code. It seems to me that it is not really of particularly high severity, though it is an easy correctness win to fix it.

Summary

PET turns per-key cutoff factors into an additive attention bias with log(clamp(factor, 1e-15)). The clamp avoids log(0), but it maps a key with cutoff factor exactly 0 to a bias of log(1e-15) ≈ -34.54 instead of -inf. Such a key therefore keeps attention weight exp(logit - 34.54) / Z rather than exactly zero — it is strongly penalised, not excluded.

The invariant that should hold: a key whose cutoff factor is 0 contributes nothing to any query's output. It does not.

Location

src/metatrain/pet/modules/transformer.py, AttentionBlock.forward (lines 109–110 at main@2618bb0):

attn_weights = torch.clamp(cutoff_factors[:, None, :, :], self.epsilon)  # self.epsilon = 1e-15
attn_weights = torch.log(attn_weights)

attn_weights is then passed as the additive attn_mask to manual_attention (line 113) and scaled_dot_product_attention (lines 133/143).

When does a cutoff factor become exactly 0?

Two ways, and the second is the one that makes this a locality bug rather than a batching artifact:

1. Padding keys. Batching pads each node's neighbour list to the batch maximum; padded slots are set to factor 0 (transformer.py:537).

2. Real edges at the cutoff. cutoff_func_bump (src/metatrain/pet/modules/utilities.py:21–22) clamps its argument to 1 - 1e-6 and then evaluates tanh(1 / tan(pi * x)). tanh saturates to -1 in floating point well before the boundary, so real neighbours in the outer ~0.02 Å (fp32) of their pair cutoff get factor exactly 0 while still being in the kept edge set (the adaptive scheme keeps every edge with r <= pair_cutoff). With the adaptive cutoff (PET-MAD targets ~8 neighbours) that shell sits, by construction, in a populated region.

   Measured on a rattled 64-atom Si cell (fp32): 14 real edges with factor exactly 0, hitting 13 of 64 nodes.

Note also that padded keys are not harmless "zero tokens": the edge embedder and RMSNorm carry biases, so a padded slot reaches the attention as a full-magnitude token (measured token norm 7.24 vs 7.12 for real edges). The -34.54 bias is the only thing holding it out.

Minimal reproducer (no checkpoint)

import torch
from metatrain.pet.modules.transformer import AttentionBlock

torch.manual_seed(0)
block = AttentionBlock(total_dim=8, num_heads=2, temperature=1.0).double()

x = torch.randn(1, 3, 8, dtype=torch.float64)
cutoffs = torch.ones(1, 3, 3, dtype=torch.float64)
cutoffs[:, :, 2] = 0.0                      # key 2 is EXCLUDED (cutoff factor 0)

out_a = block(x, cutoffs, use_manual_attention=True)
x2 = x.clone(); x2[:, 2, :] += 1.0          # perturb only the excluded key's token
out_b = block(x2, cutoffs, use_manual_attention=True)

# rows 0 and 1 are queries that did not change; an excluded key cannot affect them
print((out_a[:, :2] - out_b[:, :2]).abs().max().item())   # main: ~1e-15, not 0
assert torch.equal(out_a[:, :2], out_b[:, :2])            # fails on main

The exclusion is not exact, so perturbing an excluded key changes a real query's output. In the adversarial regime the effect is total — raw logits [0, 100] with cutoff factors [1, 0] put weight 1.0000 on the excluded key under the released code (0.0000 with a hard mask).

Measured impact

On the public PET-MAD-xs checkpoint (lab-cosmo/pet-mad), released code vs. a hard mask, same weights:

quantity	fp32 (production)	fp64
leaked attention weight (per row)	~5e-15 (up to ~4e-13 batched)	same
ΔE vs hard mask	exactly 0	~1e-12 meV/atom
max ΔF vs hard mask	exactly 0	~4e-12 meV/Å
batch invariance (structure alone vs batched)	released 9.5e-4 meV, hard mask exactly 0	released 5e-12 meV, hard mask 0

Why so small: the leak is exp(gap - 34.54), where gap is the excluded key's raw logit minus the log-sum-exp of the surviving keys. PET-MAD's trained logits are mild (max |logit| ≈ 11.7, observed max gap ≈ 1.2–4.8), so the leak is ~1e-15. The center token is hardcoded to factor 1 and is a key in every row (transformer.py:531–536), which bounds gap ≤ 2·max|logit| and caps even an adversarial arrangement of these weights at leak ≤ ~1e-5.

Error scales linearly with the clamp floor (verified over 13 decades); reaching ~1 meV/Å of force error would need the floor at ~1e-4, i.e. ~26 nats more logit gap than PET-MAD has.

So the only observable consequence for a trained model is that batching perturbs the last bit of the energy (~1 ulp). After the fix, batching is bitwise invariant. There is no accuracy claim.

Proposed fix

Keep the released equation for every nonzero factor; send exact zeros to a large-negative bias:

factors = cutoff_factors[:, None, :, :]
attn_weights = torch.log(torch.clamp(factors, self.epsilon))
attn_weights = attn_weights.masked_fill(factors == 0.0, torch.finfo(attn_weights.dtype).min)

• Masking only == 0.0 (not < epsilon) leaves the bias bitwise identical for every nonzero factor → no checkpoint migration, no version bump; trained predictions are unchanged (exactly 0 in fp32, as measured).
• finfo(dtype).min rather than -inf: gives the same exact-zero weight after softmax, but a fully-excluded row degrades to uniform attention instead of NaN. (PET's center token guarantees every production row has a live key, but AttentionBlock is public and exercised standalone, and manual_attention — the double-backward path used for conservative forces — is a plain softmax that NaNs on an all--inf row.)

Performance

The released code already passes a dense float attn_mask to SDPA, so this does not change kernel/backend selection — there is no fast path to fall off. The added cost is one elementwise compare + masked_fill on an already-materialised tensor. We could not measure a reliable delta on CPU (noise floor exceeded the effect; ~1–5% on the force path, not separable from noise). A GPU benchmark at realistic system sizes should gate the claim of perf-neutrality before merge.

Note by Marcel: I'd be a bit worried about instantiating a constant in the masked_fill -- w/o torch.compile, this makes me worry about a blocking host/device transfer? If we do this fix, Someone Should Look Into This. Might be as easy as instantiating the large negative as a device tensor beforehand or something.

Scope

AttentionBlock / CartesianTransformer are imported by experimental/flashmd and experimental/flashmd_symplectic (flashmd/model.py:20), so the fix and its regression cover those too.

Suggested regression tests

1. Exact locality (fails on main): the minimal reproducer above, as an exact-equality assertion.
2. Tighten test_pet_padding: it currently asserts allclose(atol=1e-6), which passes despite the leak. The invariant is exact — assert bitwise equality of the alone-vs-batched energy.
3. All-excluded row: AttentionBlock with all-zero cutoff factors returns finite output on both attention paths.
4. Double backward with the mask (conservative forces) stays finite.

---

🤖 Filed by Claude Code on behalf of @sirmarcel.

[sirmarcel]

Currency check: still present on today's main@4aec769f (2026-07-22). pet/modules/transformer.py and pet/modules/utilities.py are byte-identical between 2618bb0 and 4aec769f (git diff empty), so AttentionBlock.forward lines 109–110 and all other cited line numbers in those two files hold unchanged on current main. (The only drift elsewhere: experimental/flashmd/model.py import is line 21 rather than 20 on current main.)