Faster prediction: length-fitted windows, cheaper calibration fit, fewer heads read

[RobbinBouwmeester]

Rebuilt on current main, which now carries the calibration split (#121), the prediction
report (#118) and the head column source (#123). Four commits; none changes what a model
predicts beyond float32 noise.

Prediction was profiled stage by stage first, and the result contradicted the intuition that
the network is the slow part. Per 10,000 peptides on a GPU box: encoding 0.8 s, the forward
pass 0.10 s, and prediction_report with a 10,541-peptide reference 13.5 s, of which about
11 s was calibration fitting because the conformal cross-fitting repeats it once per fold. On
CPU the forward pass was 2.41 s, and the head layer was not the reason: 80 heads instead of
6,543 saved 6 %, because every peptide was convolved in a 60-position window whatever its
length, and the median peptide is 14-16 residues.

The commits

Short peptides run in a window that fits them. The trunk masks its output by the true
residue count and the encoded features are bit-identical across window lengths, so only the
convolutions can carry padding into a valid position. Each reaches
(kernel - 1) // 2 * dilation positions and the reaches add, which the model now reports as
padding_reach: 4 for the shipped architecture. A window of the longest peptide in a chunk
plus that reach is therefore exact, and it was measured exact — margin +2 still changes
predictions, +4 gives 0.00e+00 over 1,500 peptides per length band. Models that pool or stride
across positions report no reach and are left alone.

Chunks are cut by length band, not batch size. Length-sorted chunks of a fixed 4,096
peptides inherit their longest member's window, so the last chunk runs thousands of ordinary
peptides in the window a handful of long ones need: 35 % of the throughput on a mixed
20,000-peptide set. A chunk is now cut as soon as its longest peptide exceeds its shortest by
more than four residues. Enforcing a minimum chunk size, so the sparse tail rides in a wider
window, measured worse and was dropped.

Batches are encoded in one pass, and the unused rolling sum is not built. Four arrays and
four tensors per peptide, collated by a DataLoader, cost several copies of a few hundred bytes
each with a lot of Python around them; writing the encoder output into batch buffers measured
1.7x faster over the feature path (0.153 against 0.089 ms per peptide). The rolling-sum matrix
is about a tenth of the encoding work and the fused trunk deletes it on the first line of its
forward; models declare whether they read it through uses_rolling_sum, so old checkpoints
keep the array they were trained with.

The calibration fit is cheaper. Head ranking no longer materialises centred copies of the
(n, 6543) matrix — the centred target sums to zero, so the covariance is a dot product per
block and the variance follows from the block's own sums. fit no longer promotes that matrix
to float64, which only doubled a 276 MB reference for no gain. RidgeCV uses its closed-form
leave-one-out route on references of at least 2,000 peptides (0.08 s against 0.84 s at 10,541,
accuracy neutral over ten setups, median MAE ratio 1.0000); smaller references keep the fold
search, where collinear columns make leave-one-out jumpy. And the spline calibration predicts
its trails only for points outside the fitted range.

Measured

CPU, 20,000 peptidoforms including encoding, before the length work and after:

threads	before	after
4	1,465 pf/s	3,115 pf/s
8	1,780 pf/s	3,702 pf/s
32	2,424 pf/s	5,260 pf/s

prediction_report on GPU: 13.54 to 8.61 s at a 10,541-peptide reference, 7.77 to 5.14 s at
5,000, 4.65 to 4.04 s at 2,000.

Equivalence

201 tests pass, including new ones asserting that a bucketed prediction equals the full
window, that a batch encodes to what __getitem__ produces item by item, and that the
rolling-sum flag leaves every other feature untouched.

Checked end to end against a per-peptide dump written before any of this work: on PXD079349
aggregate MAE is 0.1754 against 0.1756 min and coverage 0.9052 against 0.9044, with
per-peptide predictions within 0.049 min and a median difference of 0.002; on PXD081880, whose
230-peptide reference keeps the fold-based alpha search, predictions agree to 0.003 min.

length_buckets=False forces one pass in the dataset's own window, for anyone who wants the
old path.

What changed since the first version of this PR

The head-subsetting commit is gone: #123 solves that properly, by letting a calibration index
a lazy head source rather than having the caller predict fewer heads and branch. The
calibration-fit optimisations moved into calibration/multihead.py and calibration/simple.py
to follow the split.

🤖 Generated with Claude Code