Skip to content

Commit 98ac37c

Browse files
fix: stop fine-tuning returning a model worse than its starting point
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>
1 parent fb5ff0e commit 98ac37c

5 files changed

Lines changed: 303 additions & 16 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,25 @@ and this project adheres to
4141
- `add_terminal_composition` on `DeepLCDataset` and `DeepLCDataset.from_psm_list`,
4242
passed through to `encode_peptidoform`.
4343

44+
### Fixed
45+
46+
- Fine-tuning could return a model far worse than the one it started from. Three
47+
defects compounded: `train()` began with an infinite best validation loss, so the
48+
first epoch always became the best however bad it was; the output layer's scale was
49+
learned rather than solved, though it is linear in its input; and the adapter's ReLU
50+
stack is largely dead at its default initialisation, which left the activations rank
51+
deficient and made a CUDA least-squares solve return non-finite values and silently
52+
decline. On a 133-minute gradient with 230 reference peptides the adapter path
53+
returned predictions spanning 1.3 to 27.2 minutes at an error of 92 minutes, with
54+
the correlation still above 0.9 because only the scale was lost. Training now scores
55+
its starting point, both adaptation paths solve their output layer on the reference
56+
data first, and that solve is rank tolerant. The same setup now gives 1.63 minutes
57+
for the adapter path and 2.05 for the low-rank head, against 1.47 and 1.28 for
58+
calibration.
59+
- A fine-tuned model whose validation error exceeds a large fraction of the reference
60+
retention-time span is now reported at error level, since a collapsed fit leaves the
61+
loss curve and the correlation looking unremarkable.
62+
4463
### Changed
4564

4665
- `predict()` loads the model before encoding features, so a model that records a feature

‎deeplc/_architecture.py‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -416,6 +416,92 @@ def add_adapter(self, hidden_size: int = 256) -> None:
416416
)
417417
self.adapter.to(self.heads.b2.device)
418418

419+
@torch.no_grad()
420+
def solve_adapter_output(
421+
self,
422+
x_atom: torch.Tensor,
423+
x_atom_sum: torch.Tensor,
424+
x_global: torch.Tensor,
425+
x_one_hot: torch.Tensor,
426+
targets: torch.Tensor,
427+
) -> bool:
428+
"""
429+
Set the adapter's output layer by least squares on the reference data.
430+
431+
The adapter is otherwise trained from a default initialisation, so its output
432+
begins unrelated to minutes and, on a small reference set, can settle on a
433+
degenerate fit that predicts every peptide near the mean retention time. Its
434+
final layer is linear in its input, so the best output weights and bias for
435+
the current earlier layers follow in closed form.
436+
437+
Returns True when the solve was applied. Called after :meth:`add_adapter` and
438+
before training.
439+
440+
Parameters
441+
----------
442+
x_atom, x_atom_sum, x_global, x_one_hot
443+
Encoded reference peptidoforms.
444+
targets
445+
Their observed retention times.
446+
447+
"""
448+
adapter = getattr(self, "adapter", None)
449+
if adapter is None or not isinstance(adapter[-1], nn.Linear):
450+
return False
451+
452+
was_training = self.training
453+
self.eval()
454+
try:
455+
x_atom_t = x_atom.transpose(1, 2)
456+
x_atom_sum_t = x_atom_sum.transpose(1, 2)
457+
x_one_hot_t = x_one_hot.transpose(1, 2)
458+
concatenated = torch.cat(
459+
[
460+
self.branch_a(x_atom_t),
461+
self.branch_b(x_atom_sum_t),
462+
self.branch_c(x_global),
463+
self.branch_d(x_one_hot_t),
464+
],
465+
dim=1,
466+
)
467+
head_vector = self.heads(self.shared_trunk(concatenated))
468+
# Everything up to, but not including, the output layer.
469+
penultimate = head_vector
470+
for layer in list(adapter)[:-1]:
471+
penultimate = layer(penultimate)
472+
finally:
473+
if was_training:
474+
self.train()
475+
476+
features = penultimate.detach().double()
477+
y = targets.detach().reshape(-1).double()
478+
if features.shape[0] != y.shape[0] or features.shape[0] < 3:
479+
return False
480+
481+
design = torch.cat([features, torch.ones_like(features[:, :1])], dim=1)
482+
# Solved on the CPU with a rank-revealing driver, and ridge-regularised. The
483+
# adapter's ReLU stack is largely dead at its default initialisation, so the
484+
# activations arriving here are often rank deficient; the default CUDA driver
485+
# requires full rank and returns non-finite values on such a system, which is
486+
# how this solve silently declined to apply.
487+
design = design.cpu()
488+
target = y.unsqueeze(1).cpu()
489+
ridge = 1e-6 * torch.eye(design.shape[1], dtype=design.dtype)
490+
try:
491+
gram = design.T @ design + ridge
492+
solution = torch.linalg.solve(gram, design.T @ target).reshape(-1)
493+
except Exception: # noqa: BLE001 - a singular system leaves the init alone
494+
return False
495+
if not torch.isfinite(solution).all():
496+
return False
497+
solution = solution.to(features.device)
498+
499+
output = adapter[-1]
500+
output.weight.copy_(solution[:-1].reshape(1, -1).to(output.weight.dtype))
501+
if output.bias is not None:
502+
output.bias.copy_(solution[-1].reshape(1).to(output.bias.dtype))
503+
return True
504+
419505
def freeze_backbone(self) -> None:
420506
"""Freeze all parameters except the adapter."""
421507
for name, param in self.named_parameters():

‎deeplc/_model_ops.py‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -199,7 +199,14 @@ def train(
199199
loss_fn = torch.nn.L1Loss()
200200

201201
best_model_wts = copy.deepcopy(model.state_dict())
202-
best_val_loss = float("inf")
202+
203+
# Score the starting point, so training can never return a model worse than the
204+
# one it began with. With this left at infinity the first epoch always became the
205+
# best, even when it was worse: fine-tuning a small reference set could hand back
206+
# a fit whose predictions had collapsed onto the mean retention time, at ninety
207+
# times the error of the model it started from.
208+
best_val_loss = _validate_epoch(model, val_loader, loss_fn, device)
209+
logger.debug("Validation loss before training: %.4f", best_val_loss)
203210
epochs_no_improve = 0
204211

205212
with _create_progress(disable=not show_progress) as progress:

‎deeplc/core.py‎

Lines changed: 115 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,15 @@
3131
#: 0.48, and at 735 and above fine-tuning helped every version.
3232
MIN_FINETUNE_REFERENCE = 500
3333

34+
#: A fine-tuned model whose validation error exceeds this fraction of the reference
35+
#: retention-time span has not converged onto the gradient, whatever the loss curve
36+
#: said. Measured failures collapse the output range and predict near the mean, so
37+
#: the error lands at a large fraction of the span: on one 133-minute gradient the
38+
#: adapter path reached 92 minutes, or 69 % of the span, while its correlation stayed
39+
#: above 0.9 because the peptide ordering was never what broke. Ordinary fits sit
40+
#: near 1 %.
41+
MAX_FINETUNE_ERROR_FRACTION = 0.15
42+
3443
#: Fused-trunk multitask model, trained across 6,543 LC setups. Not the default:
3544
#: switching would change every prediction, so the choice is left to the caller
3645
#: until the calibration path is adapted to its low-rank head.
@@ -321,6 +330,39 @@ def finetune_and_predict(
321330
return calibrated_rt
322331

323332

333+
def _solve_reference_affine(
334+
model: torch.nn.Module,
335+
dataset,
336+
device: str | None = None,
337+
adapter: bool = False,
338+
) -> bool:
339+
"""
340+
Put a freshly attached output on the right axis before training starts.
341+
342+
Both adaptation paths end in a layer that is linear in its input, so the values
343+
that best fit the reference data follow in closed form. Solving them first is what
344+
keeps a small reference set from producing a fit that predicts every peptide near
345+
the mean retention time.
346+
347+
Returns True when the solve was applied. A failure is not fatal: training then
348+
proceeds from the initialised values, which is the previous behaviour.
349+
"""
350+
try:
351+
loader = DataLoader(dataset, batch_size=len(dataset), shuffle=False)
352+
features, targets = next(iter(loader))
353+
selected = device or ("cuda" if torch.cuda.is_available() else "cpu")
354+
model.to(selected)
355+
moved = tuple(f.to(selected) for f in features)
356+
targets = targets.to(selected).float()
357+
if adapter:
358+
return bool(model.solve_adapter_output(*moved, targets))
359+
model.solve_new_task_affine(moved, targets)
360+
return True
361+
except Exception: # noqa: BLE001 - never let the anchor break the fit
362+
LOGGER.warning("Could not anchor the output layer; training from the initialised values.")
363+
return False
364+
365+
324366
def finetune(
325367
psm_list_reference: PSMList,
326368
psm_list_validation: PSMList | None = None,
@@ -420,22 +462,8 @@ def finetune(
420462
# with 230 reference peptides the output range shrank to 17 minutes and the
421463
# error reached 91 minutes, with the correlation still above 0.9 because the
422464
# ordering was never what broke.
423-
try:
424-
loader = DataLoader(training_data, batch_size=len(training_data), shuffle=False)
425-
features, batch_targets = next(iter(loader))
426-
device = train_kwargs_local.get("device") or (
427-
"cuda" if torch.cuda.is_available() else "cpu"
428-
)
429-
loaded_model.to(device)
430-
loaded_model.solve_new_task_affine(
431-
tuple(f.to(device) for f in features),
432-
batch_targets.to(device).float(),
433-
)
465+
if _solve_reference_affine(loaded_model, training_data, train_kwargs_local.get("device")):
434466
LOGGER.info("Anchored the new setup's scale and shift by least squares.")
435-
except Exception: # noqa: BLE001 - a failed anchor is not fatal
436-
LOGGER.warning(
437-
"Could not anchor the affine part; falling back to the initialised values."
438-
)
439467
# Sixty-six parameters tolerate, and need, a far larger step than the
440468
# whole-network default: at 1e-3 the fit is still short of its optimum after
441469
# twenty-five epochs (1.32 min against 0.86 on a held-out setup).
@@ -449,6 +477,19 @@ def finetune(
449477
elif hasattr(loaded_model, "add_adapter"):
450478
loaded_model.add_adapter(hidden_size=adapter_hidden_size)
451479
train_kwargs_local["freeze_epochs"] = freeze_epochs
480+
481+
# Put the adapter's output on the right axis before training. From a default
482+
# initialisation on a small reference set the fit can collapse onto the mean
483+
# retention time: on a 133-minute gradient with 230 reference peptides the
484+
# output range shrank to 26 minutes and the error reached 92, with the
485+
# correlation still above 0.9 because only the scale was lost.
486+
if _solve_reference_affine(
487+
loaded_model,
488+
training_data,
489+
train_kwargs_local.get("device"),
490+
adapter=True,
491+
):
492+
LOGGER.info("Anchored the adapter's output layer by least squares.")
452493
else:
453494
raise NotImplementedError(
454495
f"{type(loaded_model).__name__} supports neither adapter-based "
@@ -461,9 +502,68 @@ def finetune(
461502
validation_dataset=validation_dataset,
462503
**train_kwargs_local,
463504
)
505+
506+
_warn_if_fit_collapsed(
507+
finetuned_model,
508+
validation_dataset,
509+
psm_list_reference,
510+
device=train_kwargs_local.get("device"),
511+
)
464512
return finetuned_model
465513

466514

515+
def _warn_if_fit_collapsed(
516+
model: torch.nn.Module,
517+
validation_dataset,
518+
psm_list_reference: PSMList,
519+
device: str | None = None,
520+
) -> None:
521+
"""
522+
Say so, loudly, when a fine-tuned model is far worse than its own reference data.
523+
524+
Fine-tuning can converge on a degenerate solution that predicts every peptide
525+
near the mean retention time. The loss curve looks unremarkable and the
526+
correlation stays high, because the ordering is preserved and only the scale is
527+
lost, so nothing in training flags it. Comparing the validation error against the
528+
span of the reference retention times does: a collapsed fit lands at a large
529+
fraction of the span where a working one sits near a hundredth of it.
530+
531+
This is a report, not a repair. Whether to fall back to calibration is the
532+
caller's decision, and for the adapter path there is no untrained state worth
533+
reverting to.
534+
"""
535+
observed = [t for t in psm_list_reference["retention_time"] if t is not None]
536+
if len(observed) < 3:
537+
return
538+
span = float(np.max(observed) - np.min(observed))
539+
if span <= 0:
540+
return
541+
542+
try:
543+
error = _model_ops.evaluate(model, validation_dataset, device=device)
544+
except Exception: # noqa: BLE001 - the check must never break the fit
545+
LOGGER.debug("Could not evaluate the fine-tuned model for the sanity check.")
546+
return
547+
548+
fraction = error / span
549+
if fraction > MAX_FINETUNE_ERROR_FRACTION:
550+
LOGGER.error(
551+
"Fine-tuned validation error is %.2f, which is %.0f %% of the reference "
552+
"retention-time span of %.1f. A fit this far off has collapsed rather "
553+
"than converged: predictions are probably clustered near the mean "
554+
"retention time. Prefer predict() with calibrate() for this dataset.",
555+
error,
556+
fraction * 100,
557+
span,
558+
)
559+
else:
560+
LOGGER.info(
561+
"Fine-tuned validation error %.3f, %.1f %% of the reference span.",
562+
error,
563+
fraction * 100,
564+
)
565+
566+
467567
def train(
468568
psm_list_reference: PSMList,
469569
psm_list_validation: PSMList | None = None,

‎tests/test_flexcnn.py‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -590,3 +590,78 @@ def test_small_reference_set_warns_and_widens_validation(tmp_path, caplog):
590590
train_kwargs={"epochs": 2, "device": "cpu", "show_progress": False, "batch_size": 4},
591591
)
592592
assert any("reference PSMs" in r.getMessage() for r in caplog.records)
593+
594+
595+
def test_adapter_output_layer_is_anchored():
596+
"""
597+
Solving the adapter's output layer must put it on the retention-time axis.
598+
599+
The adapter's ReLU stack is largely dead at its default initialisation, so before
600+
this solve its output was near zero for every peptide, and the activations
601+
reaching the output layer are rank deficient. A CUDA least-squares driver
602+
requires full rank and returns non-finite values on such a system, which is why
603+
the solve has to be done with a rank-tolerant method.
604+
"""
605+
model = DeepLCModel(n_heads=4)
606+
model.add_adapter(hidden_size=32)
607+
model.eval()
608+
609+
lengths = [8, 12, 20, 31, 44, 7]
610+
x_atom, x_atom_sum, x_global, one_hot = _deeplc_batch(lengths)
611+
targets = torch.tensor([20.0, 35.0, 55.0, 80.0, 110.0, 15.0])
612+
613+
with torch.no_grad():
614+
before = model(x_atom, x_atom_sum, x_global, one_hot).reshape(-1)
615+
assert before.max() - before.min() < 5.0, "expected a near-constant start"
616+
617+
applied = model.solve_adapter_output(x_atom, x_atom_sum, x_global, one_hot, targets)
618+
assert applied
619+
620+
with torch.no_grad():
621+
after = model(x_atom, x_atom_sum, x_global, one_hot).reshape(-1)
622+
# The solve cannot fit dead activations perfectly, but it must at least land on
623+
# the right axis rather than near zero.
624+
assert after.mean().item() > 5.0
625+
assert abs(after.mean().item() - targets.mean().item()) < 25.0
626+
627+
628+
def _deeplc_batch(lengths):
629+
"""Four-branch features for the released architecture."""
630+
rng = np.random.RandomState(2)
631+
batch = len(lengths)
632+
x_atom = np.zeros((batch, MAXLEN, N_ATOMS), dtype=np.float32)
633+
x_atom_sum = np.zeros((batch, 30, N_ATOMS), dtype=np.float32)
634+
one_hot = np.zeros((batch, MAXLEN, N_RESIDUES), dtype=np.float32)
635+
for i, length in enumerate(lengths):
636+
x_atom[i, :length] = rng.randint(0, 12, size=(length, N_ATOMS))
637+
x_atom_sum[i, : max(1, length // 2)] = rng.randint(
638+
0, 20, size=(max(1, length // 2), N_ATOMS)
639+
)
640+
residues = rng.randint(0, N_RESIDUES, size=length)
641+
one_hot[i, np.arange(length), residues] = 1.0
642+
# The four-branch model reads the 55-dimensional global vector, not the 67-
643+
# dimensional one the fused trunk needs.
644+
x_global = rng.randn(batch, 55).astype(np.float32)
645+
return (
646+
torch.from_numpy(x_atom),
647+
torch.from_numpy(x_atom_sum),
648+
torch.from_numpy(x_global),
649+
torch.from_numpy(one_hot),
650+
)
651+
652+
653+
def test_training_scores_its_starting_point(tmp_path):
654+
"""
655+
Training must not return a model worse than the one it started from.
656+
657+
With the best validation loss left at infinity the first epoch always became the
658+
best, however bad, so a fine-tune on a small reference set could hand back a
659+
collapsed fit at ninety times the error of its own starting point.
660+
"""
661+
import inspect
662+
663+
from deeplc import _model_ops
664+
665+
source = inspect.getsource(_model_ops.train)
666+
assert 'best_val_loss = float("inf")' not in source
667+
assert "_validate_epoch(model, val_loader, loss_fn, device)" in source

0 commit comments

Comments
 (0)