Matrix-Preconditioned Optimizers Transfer Beyond Language Models

Abstract

Optimization algorithms transform gradient information into parameter updates for learning task- relevant representations. A family of Adam methods, which adapt each parameter independently, remains the standard choice for training neural networks. While useful, this element-wise optimization disregards the matrix structure inherent in neural-network models. Although Matrix- preconditioned optimizers exploit structural priors for layerwise updates, yielding substantial improvements across language tasks, their effectiveness on scientific data - with complex continuous inputs and non-autoregressive objectives - remains largely unexplored. In this paper, this gap is filled by benchmarking diagonal and matrix-preconditioned optimizers on peptide-spectrum retrieval in tandem mass spectrometry proteomics. Our models use a contrastive objective to embed the experimental spectra and theoretical peptides into a shared embedding space. At evaluation, each spectrum is ranked against theoretical peptides and a decoy set. We compare performance across retrieval difficulty and examine sensitivity to model size, learning rate, and batch size. On the peptide-disjoint test set, Muon, SOAP, and NorMuon consistently outperform Adam and AdamW as the candidate pool grows, while Lion provides an intermediate improvement. With 2,048 candidates per query, Muon, SOAP, NorMuon, and Lion achieve 97.88%, 97.81%, 97.74%, and 96.02% Top-1 accuracy, compared with 94.47% for Adam, which shows gains of 3.42%, 3.34%, 3.27%, and 1.55% respectively. The advantage is larger for difficult examples: for peptides of length 36-40, Muon exceeds Adam by 8.62% - with moderate increase in computational load (1.02x). Our results demonstrate that matrix-preconditioned optimization can provide substantial and increasingly pronounced retrieval gains on challenging scientific representation-learning problems, especially for mass spectrometry-based omics

Requires Python 3.11 and uv. Install the environment:

uv sync

Place the Massive-KB HDF5 files in data/proteomicsdata/:

• Massive-KB-v1_train*.h5
• Massive-KB-v1_val.h5
• Massive-KB-v1_test_sorted.h5

Weights & Biases logging runs in offline mode and does not require an account.

Training

Run the full optimizer benchmark:

uv run python run_proteomics.py \
  --config configs/run_proteomics.yaml \
  --all-optimizers

Run the learning-rate and batch-size ablation, then summarize its results:

uv run python run_proteomics_ablation_lr_batch.py --phase run
uv run python run_proteomics_ablation_lr_batch.py --phase summarize

Run the model-size ablation:

uv run python run_proteomics_ablation_model.py

Inference

After the main benchmark finishes, evaluate its best-validation weights:

uv run python infer_proteomics.py \
  --config configs/run_proteomics.yaml

Training outputs are written to outputs/run_proteomics/; inference outputs are written to outputs/run_proteomics/inference/.

License and Usage Terms

This model and associated code are released under the CC-BY-NC-ND 4.0 license and may only be used for non-commercial, academic research purposes with proper attribution. Any commercial use, sale, or other monetization of this model and its derivatives, which include models trained on outputs from the model or datasets created from the model, is prohibited and requires prior approval. If you are a commercial entity, please contact the corresponding author.