tejasjokin/ms-pred-reranker

MS3 Re-Ranker (MS3ReRanker)

Small MLP re-ranker trained via distillation from a frozen ICEBERG fragment generator + MAGMa oracle scorer.

Architecture

• Input: 13 features derived from ICEBERG top-50 fragment predictions
• Hidden: [64, 32] with LayerNorm + ReLU
• Output: scalar score (Sigmoid), trained with MSELoss against MAGMa rewards
• ~3,200 parameters

Training

• Generator: frozen ICEBERG (nist_iceberg_generate.ckpt)
• Oracle: MAGMa fragment scorer (CPU multiprocessing, 20 workers)
• Dataset: 32k samples from MSnLib splits_v2 train set
• Batch size: 256, LR: 1e-3 (Adam)

Checkpoints

File	Description
reranker_epoch01_step00025.pt	Step 25, epoch 1
reranker_epoch01_step00050.pt	Step 50, epoch 1
reranker_epoch01_step00075.pt	Step 75, epoch 1
reranker_epoch01_step00100.pt	Step 100, epoch 1 (MSE ~0.010)
reranker_epoch01_step00125.pt	Step 125, epoch 1 complete (MSE ~0.011-0.016)

Usage