(How) Learning Rates Regulate Catastrophic Overtraining

Code for the paper (How) Learning Rates Regulate Catastrophic Overtraining. arXiv:2604.13627

Installation

Requires Python ≥ 3.10 and uv:

uv sync

Data preparation

Anthropic-HH is saved once in an appropriate format and read from disk:

uv run save_hh.py --output_path data/sft/hh_messages

Running the experiments

Every experiment is a Weights & Biases sweep. For each sweep, you need to run first wandb sweep and then wandb agent (from the repo root), e.g.:

wandb sweep sweeps_colm/main/olmo12-1b.yaml   # prints a sweep ID
wandb agent <sweep-id>                        # one agent per GPU

The results are then downloadable from Weights & Biases.

Main paper

Section 3 — Low and high learning rates lead to qualitatively different models

wandb sweep sweeps_colm/main/olmo12-1b.yaml  
wandb sweep sweeps_colm/main/hubble-1b.yaml 
wandb sweep sweeps_colm/main/gemma-1b.yaml   
wandb sweep sweeps_colm/main/smollm-3.yaml 

Section 4 — Feature drift is mediated by learning rate and sharpness

wandb sweep sweeps_colm/section4/grad_step_lr.yaml           # different LRs
wandb sweep sweeps_colm/section4/grad_step_thing_olmo1.yaml  # different pretraining checkpoints
wandb sweep sweeps_colm/section4/grad_step_thing_olmo2.yaml

Section 5 — Catastrophic overtraining is caused by sharpening

Short SFT runs launched from a series of pretraining checkpoints:

wandb sweep sweeps_colm/developmental/olmo1-1b.yaml
wandb sweep sweeps_colm/developmental/olmo2-1b.yaml
wandb sweep sweeps_colm/developmental/hubble-1b.yaml
wandb sweep sweeps_colm/developmental/smollm-3.yaml

Second, the sharpness of those checkpoints, estimated by the perturbation proxy:

wandb sweep sweeps_colm/sharpness/olmo1.yaml
wandb sweep sweeps_colm/sharpness/olmo2.yaml
wandb sweep sweeps_colm/sharpness/hubble.yaml
wandb sweep sweeps_colm/sharpness/smollm3.yaml
wandb sweep sweeps_colm/sharpness/apertus.yaml 

Baseline evaluations

Benchmark scores for every base model and pretraining checkpoint, before any finetuning:

./eval_baselines.sh

Experiments in appendix

Appendix B.1 — Verifying feature drift with SAEs

Checking the alignment of MPA and SAE metrics for Gemma:

wandb sweep sweeps_colm/sae/training_with_sae_gemma.yaml

Appendix B.2 — Verifying sharpness against Hessian eigenvalues

Checks the KL sharpness proxy against true top Hessian eigenvalues on Pythia (true sharpness estimation builds on centralflows codebase):

wandb sweep sweeps_colm/sharpness_pythia/pythia_eigenvalues.yaml

Appendix B.6 — Partial pretraining experiment

Continues pretraining SmolLM3-3B from an intermediate checkpoint under two learning rate schedules that differ only in whether the LR decays. Both run on 4 GPUs.

./pretrain_smollm3_flat_only.sh        # constant LR for all 2000 steps
./pretrain_smollm3_flat_then_decay.sh  # constant for 1000, then linear decay to 0

Then measure sharpness along both trajectories:

wandb sweep sweeps_colm/partial_pretraining/sharpness.yaml

Citation

@inproceedings{
     rofin2026how,
     title={(How) Learning Rates Regulate Catastrophic Overtraining},
     author={Mark Rofin and Aditya Varre and Nicolas Flammarion},
     booktitle={Third Conference on Language Modeling},
     year={2026},
     url={https://openreview.net/forum?id=8NpiRBIijt}
}