Share one masked-marginal engine across the masked models and score the leaderboard with the wild-type fill

[Leo-T-Zang]

What this changes

The four masked language models (RNA-FM, RiNALMo, AIDO.RNA, RNAGenesis) were scored by four
near-duplicate scripts that hardcoded one masked-marginal convention. This replaces them with one
shared engine that implements the four conventions the literature defines, and moves the leaderboard
onto the one the reference implementations use.

The four fill strategies

Notation

For a variant with mutated positions $M$, write $x^{wt}$ for the wild-type sequence and $x^{mt}$ for
the variant's own sequence. At a mutated position $i \in M$, $wt_i$ is the wild-type base and $mt_i$
the mutant base. Write $x_{-i}$ for a sequence with a mask token at position $i$, and $x_{-M}$ for
one with a mask at every position in $M$. Finally $p(x_i = b \mid c)$ is the masked language model's
probability of base $b$ at position $i$ given context $c$.

Every masked-marginal score masks a mutated position, reads a log-odds between the mutant and
wild-type base there, and sums over the variant's mutations. The four conventions differ in exactly
one thing: what the model sees at the variant's OTHER mutated positions while position $i$ is
masked. All four are defined in Meier et al. 2021 (ESM-1v), supplement Appendix A.

1. wt-fill, wild-type bases at the other mutated sites

$$s_{\text{wt-fill}} = \sum_{i \in M} \Big[ \log p(x_i = mt_i \mid x^{wt}_{-i}) - \log p(x_i = wt_i \mid x^{wt}_{-i}) \Big]$$

Both terms are read from the same context, so each summand is a genuine log-odds ratio. The context
never depends on the variant, only on $i$, so a multi-mutant's score is exactly the sum of its
constituent single-mutant scores: the model is additive by construction and cannot express epistasis.
It is also by far the cheapest, needing one context per distinct mutated position of the whole assay
(1,697 across the 31 ncRNA assays) rather than one per variant.

This is what the ESM authors' examples/variant-prediction/predict.py and ProteinGym's
proteingym/baselines/esm/compute_fitness.py implement under the option name masked-marginals, and
it is the WT-LLR of the RNA zero-shot literature. The leaderboard adopts it on that provenance.

2. mask-fill, masks at the other mutated sites

$$s_{\text{mask-fill}} = \sum_{i \in M} \Big[ \log p(x_i = mt_i \mid x^{wt}_{-M}) - \log p(x_i = wt_i \mid x^{wt}_{-M}) \Big]$$

Every mutated position is masked at once and the rest of the sequence is wild type. Both terms again
share one context, so this too is a log-odds ratio. The model is told that the other mutated
positions changed but not what they changed to, so it marginalises over them rather than conditioning
on them. One context per distinct mutated-position set. This is the formula written in the ESM paper,
and its supplement calls it strategy (a).

3. mut-fill, mutant bases at the other mutated sites

$$s_{\text{mut-fill}} = \sum_{i \in M} \Big[ \log p(x_i = mt_i \mid x^{mt}_{-i}) - \log p(x_i = wt_i \mid x^{mt}_{-i}) \Big]$$

Position $i$ is masked in the variant's own sequence, so the variant's other mutations remain visible.
Both terms share a context, so it is a log-odds ratio, and it is the only one of the four that can
express epistasis, since the conditional distribution at $i$ depends on the actual genetic
background. It is also the most expensive, needing one context per (variant, position) pair
(2,458,521 across the 31 assays). Supplement strategy (c). This is the convention the previous
leaderboard used.

4. match-fill, the fill matches the allele being scored

$$s_{\text{match-fill}} = \sum_{i \in M} \Big[ \log p(x_i = mt_i \mid x^{mt}_{-i}) - \log p(x_i = wt_i \mid x^{wt}_{-i}) \Big]$$

The mutant term is read in the variant's context and the wild-type term in the wild type's, so the
two terms come from different contexts and this is not a log-odds ratio. It is a difference of two
conditionals, effectively a mutation-site-restricted difference between the variant's and the wild
type's pseudolikelihood contributions. Supplement strategy (b).

A worked example

Wild type ACGUACGUACGU with the double mutant A1C, G3U, so the variant is CCUUACGUACGU.
Scoring position 1, with # the mask token:

strategy	context at position 1	position 3 shows	terms read there
wt-fill	#CGUACGUACGU	G, wild type	$\log p(C) - \log p(A)$
mask-fill	#C#UACGUACGU	#, masked	$\log p(C) - \log p(A)$
mut-fill	#CUUACGUACGU	U, mutant	$\log p(C) - \log p(A)$
match-fill	#CUUACGUACGU and #CGUACGUACGU	mutant, then wild type	$\log p(C)$ from the first, $\log p(A)$ from the second

Position 3 is then scored the same way, and the two contributions are summed.

Two properties that follow

All four are identical on single mutants. If $|M| = 1$ there are no other mutated positions, so
$x^{mt}{-i} = x^{wt}{-i} = x^{wt}_{-M}$ and the four expressions coincide term by term. This is used
as a fixture test: on Pitt_2010_ribozyme, which is entirely single mutants, all four must return
identical scores, and they do for all eight checkpoints.

They diverge on essentially everything else here, because 99.4% of the ncRNA variants are
multi-mutants.

Why the four are computed together

Write the four half-sums

$$A = \sum_{i \in M} \log p(mt_i \mid x^{mt}_{-i}), \quad B = \sum_{i \in M} \log p(wt_i \mid x^{mt}_{-i}), \quad C = \sum_{i \in M} \log p(mt_i \mid x^{wt}_{-i}), \quad D = \sum_{i \in M} \log p(wt_i \mid x^{wt}_{-i})$$

so that $s_{\text{mut-fill}} = A - B$, $s_{\text{wt-fill}} = C - D$ and $s_{\text{match-fill}} = A - D$.
Knowing $A - B$ and $C - D$ does not determine $A - D$: adding a constant to both $A$ and $B$ leaves
mut-fill unchanged while moving match-fill. match-fill therefore cannot be recovered from the
other strategies' final scores, only from their per-position halves, which is why the engine
computes all four in one pass rather than offering them as separate runs.

Terminology hazard

The name is overloaded, and the overloading has already produced wrong numbers in this repository.
The ESM and ProteinGym code option called masked-marginals is wt-fill; the formula written in the
ESM paper is mask-fill; and wt-marginals is a different method entirely, a single unmasked
forward pass of the wild type with no masking at all, which is what the previously released RNA-FM
predictions turned out to be.

Leaderboard effect

Every masked model rises. AIDO.RNA (1.6B) 0.1824 to 0.2070, RNAGenesis 0.1498 to 0.1825, RiNALMo
0.1391 to 0.1690, RNA-FM 0.1014 to 0.1293, which passes Nucleotide Transformer, and Orthrus 0.0384 to
0.0717. The autoregressive models, NT and EVmutation are untouched, since masked marginals do not
apply to them, and their rows are byte-identical.

The five AIDO.RNA checkpoints were previously split between the leaderboard, which listed only the
1.6B, and a side table. They are five models scored on the same assays with the same convention, so
all five are now leaderboard entries and aido_rna_scaling.csv is gone.

The top three are a tie, not a ranking. AIDO.RNA (650M) 0.2163, Evo 2 (40B) 0.2120 and AIDO.RNA
(1.6B) 0.2070 span 0.0093, inside the roughly 0.01 band that 31 assays cannot resolve, which the
Notes section on that page already states. The substantive observation is not the order but that a
650M-parameter masked RNA model is level with a 40B autoregressive one.

The size series stops improving after 650M, and whether it is monotone depends on the fill strategy:
mask-fill and mut-fill increase with size, while wt-fill and match-fill put the 650M above
the 1.6B. That is now stated as prose rather than as a second table.

All four strategies are computed for every masked checkpoint and published on the leaderboard page,
so the choice can be checked rather than taken on trust. wt-fill has the highest observed macro on
8 of the 9 checkpoints, and it is the convention the leaderboard uses. The differences sit almost
entirely in the tRNA assays, which the equal-category weighting amplifies.

RNA-ERNIE is the one entry not on this convention. Its scorer runs a single unmasked forward pass
over the wild type, applies a softmax rather than a log-softmax, and sums raw probability differences,
so it is not comparable with the masked checkpoints. It needs a paddlepaddle environment we do not
have. The leaderboard page says so.

Orthrus joins the masked models. It had been excluded on the grounds that it has no mask token,
but its MLM checkpoint documents the convention directly, "positions to score should be masked
(nucleotide channels set to zero) before calling", and its published score already relied on that
operation. It now runs through the same engine as the others, and its scorer drops from 308 lines to
109 while gaining the mutant-base validation it lacked.

Implementation

fitness/baselines/masked_lm builds one deduplicated context bank covering all four strategies and
accumulates them in a single pass. The four share most of their contexts, so computing all four costs
about 19% more unique context examples than mut-fill alone (2,929,196 against 2,458,521 over the 31
assays).

Each model script is now an adapter of about 100 lines supplying an alphabet, a tokenization and one
forward pass; the four lose 1,559 lines of duplicated masking, batching and IO.

Guards, because a wrong masked score looks plausible: a tokenizer that maps a base to the unknown
token is rejected (that is the bug the released RiNALMo scores carried), every scored position is
asserted to be masked, the context-to-token offset is asserted rather than assumed, the four
strategies must cover the same variants, and the wild type is taken from RAW_CONSTRUCT_SEQ and
independently recovered by reverting each variant's mutations, with disagreement fatal.

Each prediction file gets a manifest recording the strategies, columns, alphabet, dtype, checkpoint,
counts, a hash of the scoring source and the GPU. The GPU matters: the same code and checkpoint in
bfloat16 on an L40S and an H100 differ by up to 0.2 in score and about 0.001 in per-assay Spearman.
Every number here was produced on H100s.

Verification

• --strategies mut-fill reproduces the pre-refactor predictions to 1.5e-7 (RNA-FM), 4.7e-7
  (RiNALMo), 9.2e-7 (RNAGenesis) and 7.1e-7 (AIDO.RNA, on its original L40S), with identical NaN
  masks, so the refactor changed no model's numbers.
• The mut-fill column of the new runs reproduces the previously published mut-fill leaderboard to
  within 0.0008.
• All four strategies return identical scores on Pitt_2010_ribozyme for all eight checkpoints,
  which is the fixture they must satisfy by construction.
• tests/test_masked_lm.py: 26 tests, CPU only, no checkpoint needed. Each strategy is checked
  against a per-variant reference implementation, and each adapter's alphabet, column, batching and
  dtype defaults are pinned, since those are what the released predictions were produced with.
• fitness/analyze_fill_strategies.py reproduces the sensitivity table, the category spreads and the
  bootstrap intervals from the prediction files.

Reviewing

The three commits are separable: the engine and adapters, the merge registry change, then the
leaderboard. fitness/merge_scoring_files.py keeps its previous behaviour for every existing entry;
a bare column name still means "a folder named after the model".