Fused-trunk multitask model with a low-rank per-setup head - #114
Conversation
Adds FlexCNNMultitaskModel, which merges atomic composition with a learned residue embedding in one convolutional trunk and pools over the valid length, rather than running four branches and flattening each. On the same training data, the same rank-64 head and the same 60 epochs it reaches 0.879 min test MAE against 1.263 for the four-branch backbone, improving 98.0 % of the 6,541 LC setups compared. The bundled model, trained for 200 epochs across 6,543 setups, reaches 0.823 min MAE and 0.237 min median. The head is FactorHead: a shared projection dotted with a per-setup embedding, so a setup owns rank + 2 parameters. Adding an LC setup is a 66-value fit with the encoder frozen, rather than training a head or fitting an adapter over the full head vector. Checkpoints may now describe themselves, recording architecture, constructor arguments, feature specification and target units. Loading no longer infers the architecture from tensor shapes, and predict() consults the model before encoding so a model gets the features it was trained on. This matters because the trunk fixes the width of the global feature vector: the bundled model needs the 67-dimensional form with terminal composition, and feeding it the 55-dimensional default has to fail on shape rather than silently mispredict. Bare state dicts continue to load unchanged. Target scaling is folded into the head at packaging time. Training normalises each setup's retention times, so a raw checkpoint predicts in normalised space; both transforms are affine and compose exactly, so the shipped model returns minutes with no side-car arrays. The model is not made the default, since switching would change every prediction and the calibration path has not yet been adapted to exploit the low-rank head. Verified by loading the same weights into this implementation and the training implementation and comparing outputs on identical inputs: bit-identical across peptide lengths 1 to 60. End-to-end through DeepLC's own feature extraction on 2,400 held-out test peptidoforms from six setups gives 0.686 min MAE and 0.252 min median, with Pearson r from 0.989 to 0.9998. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
CI runs `ruff format --check`, not only `ruff check`, and I had verified the latter. Two constructs in core.py and one in the new test were wrapped across lines that fit inside the 99-character limit. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
Thanks — this was a good review. All six confirmed and fixed in 5b14ad3. I reproduced each before changing anything rather than fixing from the description. 1. Full task matrix on normal predict — fixed. Confirmed the arithmetic: 6,543 columns at a million peptides is ~26 GB of output the caller discards, and I did not add a public 2. Rather than teaching 3. 4. Device — fixed. Taken from 5. 6. The padding test was vacuous — you are right, and it was worse than useless. Bundled-model test added, as suggested: asserts 6,543 tasks, 93 tests pass, |
Six issues from review, all confirmed before fixing. Normal prediction no longer evaluates every task. predict(return_matrix=False) built a column per LC setup and then discarded all but one: for 6,543 setups that is roughly 26 GB of intermediate output at a million peptides, made worse by _predict_epoch accumulating per-batch tensors before concatenating. task_idx now threads through _model_ops.predict and _predict_epoch, and the single-column path requests one column. Support is detected from the forward signature, so models without it are unaffected. save_model produced a checkpoint load_model could not read. It wrote a bare state dict for every model, and the loader then fell back to inferring the architecture, failing with KeyError: 'heads.b2'. Serialisation is now the model's own business: FlexCNNMultitaskModel.describe() returns architecture, constructor arguments, feature spec, target units, task names and weights, and save_model uses it when present. Confirmed by a test that goes through the public save_model rather than hand-building the dictionary. finetune() called add_adapter() unconditionally, which this architecture does not implement, so it raised AttributeError. It now raises NotImplementedError naming the model and pointing at predict() with calibrate(). Fitting the per-setup parameters directly is the better answer and remains follow-up work. predict() loaded the model before reading the requested device, so asking for CPU on a CUDA machine placed the checkpoint on the GPU first and could fail with an out-of-memory error for a caller who wanted neither. The device is now taken from predict_kwargs before loading. Only two of the recorded feature flags were honoured. padding_length is now threaded through DeepLCDataset to encode_peptidoform, which matters here because the trunk pools rather than flattens: a mismatch changes the representation without changing any shape, so it would not raise. test_padding_does_not_change_prediction did not test padding. It cloned the same tensors and compared a prediction against itself, so it could not have caught a masking regression. It now places the same ten residues in a length-20 and a length-60 array. Adds tests for each of the above, and one against the bundled checkpoint rather than only the synthetic model, covering task count, task-name length, target units and prediction scale. That last one catches packaging drift the synthetic round-trip cannot. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…k head finetune() previously refused this architecture, pointing at calibration instead. It now fits the new setup's own rank + 2 parameters with the encoder and every pretrained setup frozen: 66 values at rank 64, against roughly 1.7 million for an adapter over a 6,543-wide head vector. That was the stated reason for the head's design and it had no public path until now. FactorHead.add_task() attaches the new setup's parameters as separate tensors rather than extra rows, so freezing the pretrained setups is a matter of requires_grad and needs no per-row gradient masking. Output collapses to one column while a new task is attached, so the training loop and predict() need no special case. Two details that are wrong in the obvious form and are documented in place: The affine part is initialised from the trained setups' own scale and shift, not from the target retention times. scale multiplies a dot product in the normalised space the model was trained in, roughly 0 to 100, so seeding it with a spread measured in minutes overshoots about thirtyfold; on a held-out setup that put the first prediction 600 minutes out. Sixty-six parameters need a much larger step than the whole-network default. At lr 1e-3 the fit is still well short after 25 epochs, 1.32 min against 0.86 on a held-out setup, so this path defaults to 0.05. finetune() also now builds its datasets from the model's recorded feature specification, as predict() already did, and loads the model once instead of twice. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Benchmarked on six LC setups from PRIDE projects submitted after the training crawl, so unseen by construction. Fine-tuning was worse than calibration below roughly 500 reference peptidoforms and better above 700, monotone across the six with no exceptions. The worst case took a 133-minute gradient with 230 reference peptides from 1.47 to 91.7 minutes. The mechanism is not overfitting. The affine part of the head collapses: the output range shrank to 17 minutes against a 133-minute gradient and every prediction piled up near the mean retention time, while the correlation stayed above 0.9 because the peptide ordering was never the problem. scale decays toward zero and shift absorbs the mean. scale and shift are linear in the prediction, so for a given embedding the best values follow in closed form. They are now solved by least squares on the reference data before training rather than learned, which both prevents the collapse and gives the embedding a sane error to descend from. On the setup above that takes the fine-tuned error from 11.4 to 2.2 minutes for the fused-trunk model, with the correlation back to 0.975. Also warns below MIN_FINETUNE_REFERENCE and widens the validation split there, since the default 10 % left 23 PSMs to early-stop against on that setup, which is noise rather than a stopping signal. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Held-out benchmark, and two things to look at before 4.1.0Benchmarked this against the two released versions on LC setups none of them was trained on. Every project is numbered above PXD079334, the highest accession the training crawl reached, so they are unseen by construction rather than by filtering a list. Six setups, one per project, the deepest run of each, across four search engines. Each version runs with its own model and its own feature code, since that is what installing it gives you. 4.0.0 and 4.0.1 ship byte-identical weights (sha256 Results, pooled over all scored observations
Paired over setups, the PR model calibrated is at 0.809× 4.0.0 calibrated, better on 83%. It also beats 4.0.0's fine-tuned result while merely calibrated, so the architecture is doing the work rather than the adaptation. 1.
|
| project | reference PSMs | 4.0.0 cal → ft | PR cal → ft |
|---|---|---|---|
| PXD081880 | 230 | 1.472 → 91.709 | 1.280 → 11.367 |
| PXD079795 | 300 | 0.373 → 0.484 | 0.229 → 0.410 |
| PXD079655 | 735 | 1.131 → 1.065 | 1.192 → 1.107 |
| PXD081924 | 1,103 | 0.439 → 0.414 | 0.364 → 0.273 |
| PXD079927 | 2,056 | 3.196 → 1.715 | 0.895 → 0.821 |
| PXD080826 | 4,133 | 0.611 → 0.578 | 0.481 → 0.381 |
Monotone, no exceptions: worse below ~500 reference peptidoforms, better above ~700.
The mechanism is not overfitting. On PXD081880 the fine-tuned output range collapsed to 17 minutes against a 133-minute gradient, with every prediction near the mean RT, while r stayed at 0.92 — the peptide ordering was never the problem, the scale was. scale decays toward zero and shift absorbs the mean.
Since scale and shift are linear in the prediction, fb5ff0e solves them by least squares on the reference data before training instead of learning them. On that setup the fused-trunk model goes from 11.4 to 2.2 minutes with r back to 0.975. It also warns below MIN_FINETUNE_REFERENCE = 500 and widens the validation split there, because the default 10% left 23 PSMs to early-stop against, which is noise.
I have only applied the least-squares anchor to the low-rank head. The adapter path is still exposed to the same failure — 4.0.0 went to 91.7 min and I have not touched that code. Worth deciding whether 4.1.0 should guard it too.
2. 4.0.1 looks like a regression on modified peptides
Reproduced across three independent designs (30 setups, 18 setups, 6 setups): median ratio 1.017 against 4.0.0, better on one of six setups. Worst on the TMT phosphopeptide project, 3.196 → 4.686 min calibrated, +47%.
Unmodified peptides are unaffected — the two agree to within 0.005 min. Comparing feature output directly, matrix_global is identical between the versions for unmodified peptides and differs for modified ones, same length so nothing raises. Combined with byte-identical weights, that means the corrected features are being fed to a model trained on the uncorrected ones.
The correction itself looks right; what is missing is a retrain. I would treat this as needing an independent check before acting — it is one benchmark, six setups, one seed — but if it holds it argues for retraining the bundled model against the corrected features rather than shipping the correction alone.
Caveats I would not drop
- Six setups, one split, one seed. The per-setup table is the robust part; the pooled decimals are not.
- Fine-tuning is unstable run to run. PXD081924 gave 3.30 min in one run and 0.27 in another, same data, same 30 epochs, no seed set.
- Two datasets excluded, for stated reasons. PXD082486 gives r ≈ 0.35 for every version, and reading unmodified peptides straight from its source SQLite with no conversion still gives r = 0.565, so its retention times are not predictable by any model. PXD080314 is peptidomics of non-specific cleavage products and sits near 5.8 min for everything.
- Coverage is bounded by what converted. 11 of 20 downloaded files produced no tasks, including every mzIdentML file and a DIA-NN Parquet report, so this covers mzTab, DIA-NN, FragPipe and MaxQuant only.
- Effort is not symmetric. Fine-tuning fits ~1.7 M parameters for the released versions and 66 for this one, at the same 30 epochs.
96 tests pass, ruff check and ruff format --check clean.
Three defects compounded, each hiding the next. train() started with best_val_loss at infinity, so the first epoch always became the best however bad it was, and the state training began from was discarded. It now scores that starting point first, which means a fit can never be returned that is worse than what it started with. The output layer's scale was learned rather than solved. Both adaptation paths end in a layer linear in its input, so for given earlier layers the best values follow in closed form; they are now solved on the reference data before training. The adapter's ReLU stack is largely dead at PyTorch's default initialisation: before training its output was 0.0 to 0.1 for every peptide. That leaves the activations reaching the output layer rank deficient, and torch.linalg.lstsq on CUDA requires full rank, so the solve above returned non-finite values and declined by design without saying so. It now uses a ridge-regularised normal equation on the CPU, which tolerates rank deficiency. Measured on the held-out setup that failed worst, a 133-minute gradient with 230 reference peptides. The adapter path went from 92.1 minutes to 1.63 and the low-rank head from 11.4 to 2.05, against 1.47 and 1.28 for calibration. Before the fix the adapter predicted a 26-minute range for a 133-minute gradient while keeping a correlation above 0.9, because the peptide ordering was never what broke. Also reports, at error level, a fit whose validation error is a large fraction of the reference retention-time span. A collapsed fit leaves both the loss curve and the correlation looking unremarkable, so nothing else flags it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The add_terminal_composition and padding_length descriptions were spliced into each other in both DeepLCDataset.__init__ and from_psm_list, so each parameter carried half of the other one's text and two "Default is" lines. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…e 4.0.1 df4f3f6 corrected where modification deltas land in the positional block: rows are laid out as sorted(positions), so raw indexing put an N-terminal delta four residues from the C-terminus, and an if/elif reached only one row where the base residue occupied two. That correction shipped in 4.0.1 and changes the encoding of every modified peptidoform, so any model trained before it now sees input it was not trained on. legacy_positional_deltas reproduces the old placement for those models. It threads through encode_peptidoform, DeepLCDataset and feature_spec, defaults to False, and is a no-op on unmodified peptidoforms. Verified against a real v4.0.0 checkout: 4,760 feature arrays from 1,190 peptidoforms, covering lengths 2 to 70, every modified position, terminal and both-terminal modifications, are bit-identical with the flag set. The corrected path differs from v4.0.0 on 525 of them, all matrix_global, confirming the flag is the only difference. End to end, IM2Deep 2.0.2 passing the flag through its single DeepLCDataset.from_psm_list call reproduces its v4.0.0 CCS predictions exactly (max |delta| 0 over nine peptidoforms); without it, modified peptides move up to 7.6 A^2. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…d on 4.0.1 corrected where modification deltas land in the positional block, which changed the encoding of every modified peptidoform. Every model released up to that point was trained on the old encoding, so from 4.0.1 they were predicting modified peptides from input they had never seen, silently and with no error. The layout a model expects is now resolved from the specification it carries, in one helper used by predict() and finetune(). A checkpoint with no specification was written before 4.1.0, hence before the correction, so it gets the old placement; all five bundled checkpoints are bare state dicts and fall in that class. A checkpoint that records a specification is read literally, and every checkpoint this version writes records the encoding it used. The bundled FlexCNN model now declares legacy_positional_deltas: false, which is what it was trained with. DeepLCDataset defaults to the old placement, because a dataset exists to feed a model and the models in the field expect it. encode_peptidoform keeps defaulting to the corrected placement, since its job is correct featurisation. train() passes the corrected placement explicitly, so new models are unaffected. Verified against a v4.0.0 checkout, both without changing anything downstream: the default bundled RT model reproduces v4.0.0 predictions exactly on nine peptidoforms including phospho, oxidation and TMT, and pristine IM2Deep 2.0.2 reproduces its v4.0.0 CCS predictions exactly, where under 4.0.1 they had moved by up to 7.6 A^2. Re-saving the FlexCNN checkpoint with its declaration leaves its own predictions bit-identical. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
1 participant
What this adds
FlexCNNMultitaskModel, an alternative multitask backbone, plus the model trained on it and the loading machinery it needs.The existing
MultitaskDeepLCModelruns four independent branches over four feature arrays and concatenates the flattened results (1,484 values). This one passes atom counts through a pointwise stem, concatenates a learned residue embedding, runs a convolutional stack, and pools over the valid length. Pooling instead of flattening also makes it length-agnostic.Accuracy
Same training data, same rank-64 head, same 60 epochs, so only the encoder differs:
deeplc_fixed)flex_cnn_fixed)Better on 98.0 % of the 6,541 LC setups with more than 20 test observations, median per-setup ratio 0.71. The bundled model is the same architecture trained 200 epochs: 0.8233 min MAE, 0.2370 min median, 2.4464 min p95.
The head
FactorHeadcomputespred[:, j] = (proj(trunk) · embedding[j]) * scale[j] + shift[j]. The projection is shared, so a setup ownsrank + 2parameters — 66 at rank 64. Adding an LC setup is a 66-value fit with the encoder frozen, against roughly 1.7 M parameters for an adapter MLP over a 6,543-wide head vector.I have not wired this into
calibrate(). The model works through the existing path (predict the full matrix, pick the best-correlating head, fit a spline), and adapting calibration to fit the 66 parameters directly is a natural follow-up but a separate change.Self-describing checkpoints
A checkpoint may now record
architecture,encoder_kwargs,head_kwargs,feature_specandtarget_units, soload_modelbuilds the right class instead of inferring it from tensor shapes. Bare state dicts still load unchanged, and there is a regression test for that.This is load-bearing rather than cosmetic. The trunk fixes the width of the global feature vector: this model needs the 67-dimensional form produced with
add_terminal_composition=True, and the 55-dimensional default must fail on shape rather than silently mispredict. Sopredict()now loads the model before encoding features, andDeepLCDatasetgainedadd_terminal_composition.Units
Training normalises each setup's retention times, so a raw checkpoint predicts in normalised space. Both transforms are affine and compose exactly, so packaging folds the normalisation into
scaleandshift; the shipped model returns minutes with no side-car arrays. Worth noting the existingmultitask_model.ptreturns roughly 0–1.3 on the same peptides and needs calibration to become minutes, whereas this one does not.Not the default
DEFAULT_MODELis unchanged. The new file is exposed asdeeplc.core.FLEXCNN_MULTITASK_MODEL. Switching the default would change every prediction and should be a deliberate decision, ideally after calibration exploits the low-rank head.Testing
tests/test_flexcnn.py, 18 tests, and the existing 68 still pass (86 total).ruff checkclean ondeeplc/and the new test file.Beyond the unit tests, two checks that seemed worth doing before trusting the port:
Numerical equivalence. The same weights loaded into this implementation and into the training implementation, run on identical inputs. Differences were exactly zero across peptide lengths 1, 2, 3, 8–40, and 60, and for a batch mixing 1 to 60. The port reorganises the module tree, drops unused options and takes one-hot input where the original took integer indices, so this was measured rather than assumed.
End to end through DeepLC's own extractor, on 2,400 held-out test peptidoforms from six well-populated setups:
Consistent with the run's own figures, which it would not be if the feature contract, the one-hot to index conversion or the folded scaling were wrong.
Things a reviewer should push back on
multitask_model.ptat 23 MB, but if models should move out of git this is the moment to say so.x_atom_sumis accepted and ignored, to keep the signature interchangeable withMultitaskDeepLCModel. A cleaner signature would break that interchangeability; I picked compatibility, but it is a judgement call.🤖 Generated with Claude Code