Trainer optimizes against AF-permuted labels (aatype never converted to MPNN)

[maraxen]

Found during the campaign control-knob audit (task_id: 260715_aminx-campaign-control-knob-audit, PR #108). Confirmed empirically, not by reading. Filed rather than fixed because the fix has consequences beyond the code — see "Decision needed".

The bug

aatype is AF-ordered (ARNDCQEGHILKMFPSTWYVX); the model's token space is MPNN (ACDEFGHIKLMNPQRSTVWYX). The training path never converts.

Traced, every hop, no conversion anywhere:

• training/dataloading/preprocess.py:126 persists protein.aatype raw, under the unsuffixed key "aatype".
• trainer.py:780 / :819 / :909 bind batch.aatype / val_batch.aatype / test_batch.aatype to train_step's sequence param.
• trainer.py:301 declares sequence: jax.Array; its docstring :326 reads "Target sequence (integer labels)" — names no alphabet.
• trainer.py:435 — cross_entropy_loss(logits, seq, ...), where logits come from m.decoder.call_conditional, i.e. MPNN. sequence is the CE target.
• trainer.py:375, :635 — one_hot(seq, 21) as the decoder embedding input, not just the label.
• trainer.py:667, :689 — eval_step / batch_metrics. So sequence_recovery_accuracy and perplexity are also computed against permuted labels — which is why nothing ever looked wrong.
• train_diffusion.py:44 — same pattern.

The structural fact: af_to_mpnn( appears in exactly one place in all of src/ — aa_convert.py:97, inside string_to_protein_sequence, which takes a str. No array-valued path in aminx converts anywhere.

Empirical confirmation

The training ingress is proxide.ops.dataset.create_protein_dataset (trainer.py:17), which never touches io/parsing/dispatch.py. Ran it directly on tests/data/1ubq.pdb:

AF   decode: MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQ   <- 1UBQ's real sequence
MPNN decode: PGLQYNTMTINTLTMHYHRSETLHDYNANLGENHILRREG   <- garbage

Against real proteinmpnn_v_48_020 weights, native recovery is 0.592 via af_to_mpnn vs 0.118 raw.

Severity depends on initialization

• Fine-tuning from an MPNN checkpoint (trainer.py:159/:187 both call load_model): straightforwardly corrupt — gradient descent is pushed to relearn a permutation the weights already encode correctly.
• Training from scratch: the model self-consistently learns AF order. No loss anomaly, no symptom at training time — and the checkpoint silently becomes AF-native while the whole inference stack assumes MPNN. Arguably the nastier case.

Scope

• No impact on inference with the bundled pretrained checkpoints. sample() is clean; those weights are MPNN-native by construction.
• No impact on tev_design, which never trains.
• Any model aminx has trained or fine-tuned is suspect, and so is every metric reported for it.

Decision needed (why this is filed, not fixed)

The code fix is mechanical (af_to_mpnn at the sites above). The consequences are not:

1. Do any aminx-trained checkpoints exist? If so they are AF-native and need retraining or a permuted output head.
2. Is preprocess.py:126's persisted "aatype" key rewritten (invalidating existing array_record datasets) or converted at read time?
3. Historical training metrics should be marked unreliable.

Guard that would have caught it

PR #108 adds tests/utils/test_alphabet_boundary.py — a native-recovery differential against real weights (0.592 vs 0.118). It keys on a distributional comparison to ground truth, not on values: a value check is provably impossible here since every int 0–20 is legal in both alphabets. Extending it to cover one train step is the natural gate for this issue.

[maraxen]

Triage: LATENT, low priority. Confirmed with the maintainer — the trainer has never been run. Nobody has trained or fine-tuned a model with it.

That resolves every open question in the "Decision needed" section:

1. No aminx-trained checkpoints exist, so nothing is corrupted and nothing needs retraining or a permuted output head.
2. The preprocess.py:126 persisted-key question is moot — there are no array_record training datasets to invalidate. Whichever fix is chosen (rewrite the key vs. convert at read) can be decided on its merits when someone picks this up, not under migration pressure.
3. No historical training metrics to mark unreliable.

So this is a real bug in a feature that has never executed. The bug report above stands as written — the tracing is unchanged and the empirical confirmation (create_protein_dataset yields AF; trainer.py:435 uses it as the CE target against MPNN logits) still holds.

What matters now is that it must not fire the first time someone uses the trainer. The natural gate is already in the repo: extend tests/utils/test_alphabet_boundary.py (PR #108) to cover one train step, asserting native recovery beats chance through the real training path. That would make this impossible to ship rather than requiring someone to remember it.