# AxoMEME 2.0 — viral fine-tuned model, artifacts and scripts **Owner:** Steven Weaver · **Host:** silverback · **Path:** `/archive/sb-data/shares/web/web/stevenweaver/axomeme-2.0-viral` **URL:** Everything needed to run, test or further fine-tune the model, in one place. Verify what you downloaded against `SHA256SUMS` before using it — the same weights have circulated under several filenames, so the hash is the only reliable identity. ## Files | file | sha256 (first 12) | what it is | | :--- | :--- | :--- | | `model/axomeme_2.0_viral_finetuned.onnx` | `3e06b591a060` | **The one DataMonkey3 ships.** Model-only graph, opset 17, eval mode (dropout off). | | `model/axomeme_2.0.onnx` | `5b970f0bff2a` | Sergei's released TOGA-mammal 2.0, same architecture. For comparison, not for serving. | | `model/axomeme_2.0_viral_finetuned.pt` | `1f860832894d` | The PyTorch checkpoint the ONNX above was exported from. Start here to fine-tune further. | | `scripts/predict_regression_nexus.py` | — | Upstream inference driver. **Read the known issues below before trusting its output.** | | `scripts/train_transformer_selection.py` | — | Model class and tokenizers. `predict_*` imports from it. | | `scripts/requirements.txt` | — | torch, numpy, pandas, scipy, biopython, scikit-learn. | `axomeme_2.0_viral_finetuned.pt` is byte-identical to what was also circulating as `recon_2.0_finetuned.pt`. Same file, two names. ## The ONNX graph is the model only Five already-computed tensors in, five out: ``` msa_codons int64 [B, N, 1] codon token ids, TCAG order (see issue 1) msa_aas int64 [B, N, 1] amino-acid token ids dist_matrix float [B, N, N] patristic distances mds_coords float [B, N, 4] MDS of the distance matrix — an INPUT, not computed in-graph padding_mask bool [B, N] TRUE = padded -> lrt, alpha, beta_neg, beta_pos, p_neg float [B] ``` `torch.linalg.eigh` has no ONNX lowering, so MDS cannot go in the graph. Anything consuming this must reproduce the preprocessing itself: newick parse, patristic distances, Max-PD taxon selection, tokenisation, MDS. DataMonkey3 has a JS port verified against these scripts to Pearson r = 1.000000 over 2,195 sites. The rate heads are `softplus`, i.e. **log1p(rate)** — apply `expm1` to `alpha` / `beta_pos` before reading them as rates. `p_neg` is already a probability; report `1 - p_neg`. ## Known issues in the shipped driver All three are one-to-four-line fixes and all three change results. Fix them before generating anything you intend to trust. **1. The codon tokenizer disagrees with training on 63 of 64 codons.** `predict_regression_nexus.py` defines its own 60-codon *alphabetical* `CODON_LIST` (lines 49-55), then redefines `get_codon_token` (line 74) without redefining `CODON_TO_IDX`, and imports only three non-tokenizer names from the training module. `TTA` and all three stop codons are absent from its list entirely and collapse to "unknown". Training uses 64 codons in TCAG order. `ATG` is 35 at training and 14 at inference. Fix: import the tokenizers from `train_transformer_selection.py` instead of redefining them. **2. Inference is non-deterministic — dropout is left on.** Line 1364 calls `model.train()` to reach the raw-ordinal-logits branch, which also enables dropout, and it is never disabled. Two runs of the same input differ by up to **0.41** in `predicted_lrt`, against calling gates of 3.12 and 4.45. This also makes the driver disagree with the ONNX export by construction, since the export is eval mode. Fix: ```python model.train() for m in model.modules(): if isinstance(m, torch.nn.Dropout): m.eval() ``` **3. No clamp on negative distances.** Line 955 computes `log((node_count + 1.0) / (dist + 0.1))` unguarded. DataMonkey's neighbour-joining emits negative branch lengths, so this raises `ValueError: expected a positive input` on **13 of 283** real submissions (4.6%), worst case −114.9. The *training* pipeline clamps distances ≥ 0; inference should too. ## What the model is for It **ranks** sites by how MEME-like their signal looks. It is not calibrated to MEME's LRT scale: across 12 real submissions and 662 variable sites the predicted LRT reached the p ≤ 0.10 gate once and the p ≤ 0.05 gate never — including on an alignment where MEME reports 17 significant sites. The reported metric is Spearman rank correlation (0.481 on held-out viral families), which is a claim about *order*, not about significance. Use percentile or z-score calling. The driver's `pvalue` default will report nothing on real data. ## Fine-tuning further Start from `axomeme_2.0_viral_finetuned.pt`. It is architecture self-describing — `predict_regression_nexus.py` auto-detects `embed_dim=128`, `num_layers=4`, `num_heads=4`, `window_size=1`, `num_streams=4`, `pure_coral=True`. Instantiating the class directly requires `num_heads=4`; it is not the default. `torch.load(..., weights_only=False)`; `checkpoint['phase2']` and `checkpoint['val']` carry provenance. **Training data is not here and must not be put here.** It is real DataMonkey submissions — unpublished research — and stays on the cluster. ## Open question for whoever picks this up If any part of the fine-tune ran on driver-tokenised inputs (issue 1), the checkpoint learned the wrong codon mapping and needs retraining rather than a corrected script. That has not been confirmed either way.