perf(inference): drop training-only targets the inference forward never reads

[mooreneural]

At inference the featurizer builds two tensors that Nesso1.forward never reads: disto_target and token_to_rep_atom. Both are training-supervision artifacts inherited from the Boltz featurizer nesso was adapted from: disto_target is the ground-truth distogram label for the distogram head's cross-entropy loss (an O(N^2 * num_bins) one-hot built with cdist + one_hot), and token_to_rep_atom is the token-to-representative-atom gather (O(N_tok * M_atoms) one-hot) for the structure/confidence heads. The codebase is inference-only, so this removes them outright, along with the now-dead _compute_disto_target helper.

They were previously constructed, collated, and copied to the device every batch, then ignored. Removing them is output-identical.

(Per @shenoynikhil's review: dropped the earlier compute_training_targets flag and just remove them, since the repo is inference-only.)

What changed

• nesso/data/featurizer.py: process_token_features and process_atom_features no longer build disto_target or token_to_rep_atom, and _compute_disto_target is removed. Net ~40 fewer lines; inference.py is unchanged from main.
• tests/test_inference_features.py: pins that the featurizer omits both keys, and that the forward (trunk and affinity head) is bit-identical whether or not those tensors are injected into the batch, so the model provably never consumed them. Ligand-only input keeps it CCD-free for CI.

Correctness

Neither key is referenced anywhere in nesso/model/ (grep). The removed code draws no RNG and touches no other feature, so every tensor the model reads is byte-identical. The new tests confirm forward output is torch.equal-identical with vs without the tensors present, on both the trunk and the affinity head. Full suite: 23 passed, 7 skipped (the 7 skips are the existing CCD-gated protein tests).

Impact (measured, RTX 5080, N=740 tokens)

• input-batch GPU memory: 155.1 → 16.8 MiB (~9x smaller; these are exact tensor sizes)
• host-to-device copy: ~9.8 → ~1.3 ms per batch (~7x faster)
• disto_target alone is exactly 64 MiB at N=512 and 144 MiB at N=768, all eliminated

Featurization is also ~1.3-1.6x faster per record, but that runs in DataLoader workers overlapped with GPU compute, so treat it as secondary; the memory and transfer numbers are the reliable ones.

[shenoynikhil]

Good catch @mooreneural! Can we actually just remove the two features because the current codebase is inference-only?

[mooreneural]

Makes sense, since it's inference-only there's no reason to keep the flag. I'll remove both tensors outright and drop compute_training_targets. Updating the PR now.

[mooreneural]

Thanks @shenoynikhil! Agreed, done. The featurizer change is now a net ~40-line removal (no more disto_target / token_to_rep_atom, plus the dead _compute_disto_target helper), and dropping the compute_training_targets flag means inference.py is unchanged from main. The rest of the diff (~140 lines) is the new regression test.

That test keeps the correctness proof: the forward, on both the trunk and the affinity head, is bit-identical (torch.equal) whether or not the two tensors are injected into the batch, so the model provably never read them. It's a bit fuller than the change strictly needs, though. Happy to trim it to just asserting the featurizer omits both keys if you'd rather keep it lean.

I also dropped the benchmark script, since its full-vs-lean comparison relied on the flag. Happy to add a standalone one back if that's useful.