Source-linked AI summary

Post-Training Science for Supervised Fine-Tuning

Charles O'Neill, Mudith Jayasekara, Harry Partridge

arXiv:2609.01244v1cs.LGcs.CL

TL;DR

Supervised fine-tuning still relies on per-job heuristics, so the paper measures its key decisions systematically across models, tasks, architectures, and training methods. Controlled one-variable-at-a-time sweeps identify transferable defaults and boundaries, including a flat LoRA learning-rate rule, rank and alpha settings, conditional validity of validation loss, scaling patterns, and an approximately two-epoch ceiling. The conclusions are directional and task-scoped because the datasets are judge-optimised and several experimental axes remain limited.

  • Problem

    SFT decisions are commonly inherited from defaults, tuned on other models, or rediscovered through ad-hoc sweeps for each job.

  • Method

    The paper performs controlled one-variable-at-a-time sweeps across two model families, dense and mixture-of-experts models, four customer datasets, LoRA, and full fine-tuning.

  • Results

    The optimal LoRA learning rate is 10−3 across 0.6–32B and both families; rank 64 with α = 32 is the default, validation loss transfers only within fixed recipes, and useful training generally ends near two epochs.

  • Takeaways & Limitations

    The study supports reusable SFT defaults within its tested scope while requiring downstream checks for cross-family selection and monitoring capability erosion with additional epochs.

  • Takeaways & Limitations

    The results use four judge-optimised datasets, leave several comparisons open, and treat directions rather than point estimates as robust because each family has few model sizes.

Abstract

from arXiv · show

Every supervised fine-tuning run forces the same chain of decisions, such as learning rate, batch size, LoRA or full fine-tuning, how many epochs, which optimiser, and what data to feed the model. Each of these is typically rediscovered from scratch for every new model and dataset. Here we measure them under one instrument: a sweep that varies one lever at a time, and spans dense and mixture-of-experts models in two families (Qwen3 and Llama), on four real-world customer SFT datasets, for both LoRA and full fine-tuning. These datasets give a controlled testbed: each task carries an evaluation built with the customer, and its training data is produced by iterative supervised fine-tuning that refines model outputs until they pass that evaluation, so the supervised target is internally consistent and the task judge we report against is the criterion the data was built to satisfy. We ask how the optimal learning rate and batch size move with model scale, family, and data, and whether one selection rule transfers across them; what LoRA trades against full fine-tuning, and how its rank and alpha set what the adapter can learn; whether validation loss (or other metrics, such as loss landscape flatness) faithfully ranks downstream quality; whether post-training gains scale with model size and data volume, on a model ladder extended through mixtures-of-experts to 235B parameters; how many epochs to train before general instruction-following erodes; and whether a geometry-aware optimiser improves on AdamW. Each recommendation is paired with a measure of its uncertainty.

1. Introduction

The paper replaces heuristic SFT decisions with controlled, one-variable-at-a-time measurement across models, families, fine-tuning methods, and customer tasks. It identifies transferable settings and boundaries for learning rate, adapters, validation loss, scaling, data, optimisation, and training duration.

  • Experimental scope: Controlled sweeps vary one SFT lever at a time across Qwen3 and Llama models, dense and mixture-of-experts architectures, four customer datasets, LoRA, and full fine-tuning.Data, splits, and seeds are held fixed so only the variable under study changes.
  • Controlled testbed: Customer-built evaluations and iteratively refined training examples create internally consistent targets and align the downstream judge with the criterion used to construct the data.This design makes changes in validation NLL more attributable to the lever under study rather than contradictory answers or label noise.
  • Calibration: The optimal LoRA learning rate is flat at 10−3 across 0.6–32B models and both families, roughly 33× the full fine-tuning optimum.The rule transfers to a held-out 30B mixture-of-experts model and preserves the monotone size trend at 80B and 235B.
  • Calibration: Adapter rank adds usable capacity up to about 64 before plateauing, while α = 32 is best in every cell; r = 64,α = 32 is the default.Rank 32 is cheaper and gives up at most 0.003 nats.
  • Selection metric: Validation loss ranks downstream quality within fixed model–dataset–recipe cells, with within-dataset-standardised Spearman −0.38 to −0.88, but not across model families.At matched loss, Fisher-trace flatness adds nothing reliable for in-domain selection.
  • Scaling: At-defaults loss follows a saturating power law in model size, steeper for full fine-tuning than LoRA, while tripling fresh examples lowers one-epoch LoRA loss with mostly dataset-specific returns.The model ladder extends through mixture-of-experts models to 235B total parameters.
  • Optimisation and duration: Muon’s pretraining advantage carries narrowly into SFT, while validation loss overfits past about two epochs and judged quality stops improving as instruction-following erodes.Adding fresh examples lowers loss but does not move the optimal epoch ceiling beyond approximately two.

2. Learning Rate and Batch Size

Across nine models, two families, and four customer datasets, the sweep finds transferable learning-rate rules, weak batch effects, and dataset-driven loss variation. LoRA nearly matches FullFT while using far fewer trainable parameters, but token load shifts the optimum more strongly for FullFT.

  • Experimental design: The sweep covers nine Qwen3 and Llama models, four anonymised customer datasets, fixed experimental conditions, and LoRA plus FullFT arms.Learning rate is fitted as a power law in hidden size, with loss-weighted candidate-law scoring rather than noisy per-cell argmins.
  • Learning-rate law: 10−3 is the optimal LoRA learning rate across 0.6–32B in both families, about an order of magnitude above FullFT’s near-3×10−5 optimum.The LoRA size exponent is statistically indistinguishable from zero, while the fitted multiplier and normalization trade off along a flat valley.
  • Learning-rate law: Cross-family and leave-one-dataset-out transfer is strong: a universal 10−3 LoRA rate costs under 0.01 nats relative to per-cell tuning.Cross-family prediction differs by at most 0.004 nats from the within-family law, while leave-one-dataset-out error has median 0.0 and 90th percentile ≈0.5 in log10 learning rate.
  • LoRA versus FullFT: FullFT is at or below LoRA on validation NLL in all 72 matched comparisons, but LoRA recovers a median 98% of FullFT’s improvement using 3.1–12.6% of the parameters.The comparison is matched on model, dataset, and batch at shared batch sizes {16,64}.
  • Batch size: Batch size mainly trades loss against optimiser-step cost: one-epoch wall-clock is nearly batch-invariant, while the optimal learning rate is nearly batch-invariant.The rate drifts only +0.07 dex per doubling for smaller Qwen models and not at all in Llama.
  • Dataset effects: Dataset identity explains 56–87% of final-NLL variance and 69–88% of improvement variance, while batch and learning rate each explain at most 0.07.Absolute NLL is token-confounded: token load explains R2=0.58 and 0.61 across families, and more tokens yield worse absolute NLL and smaller improvement over baseline.

3. LoRA Hyperparameters: Rank and Alpha

The adapter sweep finds a practical capacity threshold around rank 64 and an alpha optimum at 32 under the tested one-epoch, 5,000-example regime. Larger rank adds parameters with negligible validation-loss gains, while rank 32 nearly matches the default at about half the trainable-parameter fraction.

  • Rank: Rank 8 is 0.006–0.010 nats worse than the r = 64,α = 32 anchor, while rank 16 is 0.003–0.006 nats worse.
  • Rank: Rank 32 comes within 0.0009–0.0028 nats of the anchor while using about half the trainable parameters.On Qwen3-4B, the fraction is 2.25% rather than 4.40%; on Qwen3-8B, it is 1.56% rather than 3.07%.
  • Rank: The useful capacity threshold is around r = 64, not r = 128.Rank 128 improves validation loss by only 0.0004–0.0008 nats in three of four pairs and is slightly worse on Qwen3-8B/leasing.
  • Alpha: At fixed rank 64, α = 32 is best in every comparison.α = 16 hurts by 0.0022–0.0082 nats, while α = 64 hurts by 0.0009–0.0037 nats at the fixed 10^-3 learning rate.
  • Interpretation: Under standard α/r scaling, increasing rank raises parameter count while shrinking the LoRA multiplier, leaving rank 128’s remaining gain too small to justify changing the default.The added capacity outweighs the smaller multiplier from ranks 8 to 64, but not sufficiently at rank 128.
  • Scope: These adapter comparisons did not jointly sweep learning rate with rank, test a rank-stabilised LoRA variant, or run downstream evaluations.

4. Predicting Downstream Performance from Loss

Validation loss ranks downstream judged quality within fixed model, dataset, and recipe cells, but that relationship does not transfer across model families. Fisher-trace flatness follows the predicted learning-rate-to-batch trend yet adds no reliable ranking at matched loss within a recipe.

  • Within-recipe prediction: Within matched model, dataset, and recipe cells, validation loss ranks judged quality with standardised Spearman ρs from −0.38 to −0.88.This holds across both models and tuners over 896 cell-by-task pairs.
  • Base calibration: Base-model evaluation scores rise near-monotonically with scale across all four tasks, so the evaluations rank models below the quality range reached after fine-tuning.
  • Family calibration: Pooling base models weakens the loss–evaluation correlation to −0.35 because model families split the calibration.Within each family the correlation is negative on every dataset, while Llama can reach lower NLL than Qwen at matched size but score lower with judges.
  • Flatness calibration: The base Fisher trace correlates with judged quality at −0.90 pooled, but its apparent advantage largely re-expresses model scale and falls to approximately −0.15 after regressing out size.
  • Scope: The study excludes four divergent docs cells at learning rate 10^-3 because outputs exceeded the 20% judge-error gate.
  • Within-recipe prediction: At matched loss within a recipe, the Fisher trace adds no reliable ranking and its sign is unstable, including a positive relation for Llama-3.1-8B full fine-tuning.This falsifies the stated matched-loss flatness hypothesis for the converged-minimum trace at the fine-tuned scale.
  • Transfer: For transfer comparisons at genuinely matched loss, flatter cells have higher evaluation in 57–77% of pairs, although the in-domain effect remains weak.Across epoch sweeps, higher trace tracks lower IFEval capability at fixed loss but higher in-domain task score.

5. Scaling Trends in Post-Training

Post-training loss improves predictably with model size and fresh data, but scaling depends on tuning method, dataset, and how mixture-of-experts size is counted.

  • Scope and uncertainty: Scaling-law magnitudes remain uncertain because the fits use few sizes and cannot separate examples from tokens within a dataset.The design uses at least five Qwen3 sizes, but confidence intervals are wide and regime breaks are difficult to detect with few points.
  • Dense model-size trends: Loss follows a saturating power law in size, with steeper decay for FullFT than LoRA.LoRA exponents are 0.31–0.40, versus 0.40–0.51 for FullFT, with wide intervals at six points.
  • Dense model-size trends: LoRA’s FullFT-minus-LoRA gap narrows from 0.008 nats at 0.6B to 0.003 at 32B.The fitted floors track token load, and Llama remains slightly above the Qwen fit at matched parameters.
  • Mixture-of-experts scaling: The geometric mean of active and total parameters best places all three MoEs on the dense trend.Mean absolute residuals are 0.016–0.022 nats for geometric-mean placement, versus 0.021–0.073 for active or total placement.
  • Size by data: Tripling fresh examples lowers one-epoch validation NLL by about 20% on leasing and 8% on security, with a median six-curve gain of 14.4%.The fractional gain is nearly size-independent, while the dataset drives the magnitude; examples and tokens remain confounded within each dataset.

6. Optimisers

Muon reaches slightly lower loss and flatter minima than AdamW when separately tuned, but the benefit is narrow and appears mainly in retained general instruction-following rather than task-judge quality.

  • Scope: The optimiser comparison covers FullFT only, leaving Muon’s behaviour with a rank-64 adapter open.The study uses Qwen3-8B at global batch 16 for one epoch and leaves the low-rank adapter case to future work.
  • Loss and learning rate: Muon’s best validation NLL is at or below AdamW’s on every dataset, by at most 0.012 nats.Muon reaches its minimum at a learning rate about three times lower, and its frontier degrades sharply above the optimum.
  • Flatness: Muon’s per-parameter Fisher/Hessian trace is lower than AdamW’s by a factor of 1.2 to 2.8 on every dataset.The flatter minimum does not provide a reliable selection signal for in-domain quality at matched loss.
  • Downstream quality: The two optimisers are level on the task judge, with a mean difference of 0.015 within the judge’s resolution at 200 rows.This result comes from validation-NLL-selected Qwen3-8B FullFT checkpoints.
  • Downstream quality: Muon scores higher on IFEval across all four datasets, by a mean of 0.09.The comparison separates retained general instruction-following from in-domain production-judge quality.

7. Epochs

Training beyond roughly two epochs overfits validation loss and erodes general instruction-following, while fresh examples improve loss more reliably than repeated passes without justifying longer training.

  • Epoch selection: At 5,000 examples, validation NLL reaches its minimum by about two epochs and then rises monotonically to eight in every tested cell.The increase is 40–120% on harder tasks and larger still on Qwen3-1.7B.
  • Epoch selection: The judged task score holds or rises through eight epochs even as validation loss worsens.Selecting by validation loss would stop around two epochs, while the production judge rates the eight-epoch checkpoint as good or better.
  • Fresh data versus repetition: Replacing repeated 5,000-example training with fresh 10,000 examples lowers validation NLL in all 12 comparisons, by a mean of 16%.Increasing fresh data from 10,000 to 15,000 lowers NLL in all 18 comparisons, by a mean of 5.9%.
  • Fresh data versus repetition: The four-epoch run is worse than the best one- or two-epoch run in all 12 larger-data curves.The best NLL occurs at 15,000 examples and two epochs in five of six cells.
  • Capability erosion: IFEval falls from one to four epochs in 10 of 12 larger-data cells, by a mean of 0.15.Across the full grid, the cost is steepest on instruction-heavy tasks; Qwen3-8B FullFT support falls from 0.73 at two epochs to 0.34 at eight.
  • LoRA and cross-task effects: LoRA matches or beats FullFT on trained-task scores, but both methods erode IFEval similarly with more epochs.Fine-tuning also raises the trained task while lowering other task judges relative to the untuned base.
  • Scope: The 1.7B arm overfits more than the 8B models, while the 32B epoch rung remains untested because full-context evaluation exceeds single-device memory.The epoch scale grid therefore spans Qwen3-1.7B and 8B rather than the full size ladder.

8. Conclusion

The study replaces per-job post-training guesswork with controlled measurements across models, methods, datasets, and tuning levers, yielding practical defaults alongside explicit scope boundaries.

  • Contribution: Controlled one-variable sweeps across two families, dense and MoE models, LoRA and FullFT, and four customer datasets make lever effects attributable to the tested change.The customer-built data and evaluations provide internally consistent targets and a validated downstream criterion.
  • Practical defaults: The default LoRA learning rate is 10^-3, roughly 33× the FullFT optimum, while rank 64 and α = 32 remain the adapter defaults.Rank capacity plateaus around 64, and the learning-rate rule transfers to larger and held-out MoE models.
  • Selection metrics: Validation loss ranks judged quality within fixed recipe cells, but not across model families; flatness instead tracks retained general capability.The flatness relationship is clearest across optimisers rather than as an in-domain selection proxy.
  • Scaling and optimisation: At fixed defaults, loss scales as a saturating power law, LoRA’s retention gap narrows with size, and MoEs align at the geometric mean of active and total parameters.Fresh examples improve loss with mostly dataset-specific returns, while Muon’s task-judge quality is unchanged and IFEval retention is higher.
  • Caveat: Because training data is optimised to pass the customer judge, that judge is not independent of the training objective; IFEval serves as a cross-criterion check.The findings are therefore read as fit to the task as specified, rather than as independent evidence of general quality.
  • Limitations and open questions: Open boundaries include Muon with LoRA, examples versus tokens, stability-bounded learning-rate ceilings, a 32B epoch rung, and broader model or dataset coverage.With four datasets and few sizes per family, directions are treated as robust while magnitudes are reported with uncertainty.

Appendix A. Wall-clock, compute, and memory cost of the batch axis

Batch size mainly trades per-step time and memory against optimizer-step count: larger batches make each step slower, but one-epoch wall-clock remains nearly unchanged until memory limits intervene.

  • Wall-clock cost: Seconds per optimizer step grows almost linearly with global batch because each step processes approximately the global-batch examples.The timing sweep measures steady-state post-warmup step time, excluding setup, evaluation, and checkpointing.
  • Wall-clock cost: One-epoch wall-clock varies by about 10% across batch sizes because fewer steps offset the higher cost per step.At the largest batch, wall-clock is marginally lower as per-step overhead amortizes.
  • Compute efficiency: FullFT delivers roughly 1.5× the useful FLOPs per token of LoRA, producing correspondingly higher achieved FLOP/s at matched settings.The comparison follows the 6N-versus-4N model-FLOPs difference from frozen-base gradient computation.
  • Memory: Peak memory rises with microbatch size and then flattens at global batch 64, where gradient accumulation supplies additional batch.FullFT uses more peak memory than LoRA at every model size, making memory the practical batch ceiling.
  • Dataset effects: Throughput is strongly dataset-bound because sequence length differs across datasets, with support shortest and docs longest.The figure reports steady-state LoRA step time by model and dataset.

Appendix B. Dataset statistics and coverage

The appendix characterizes dataset token loads, eligible sweep coverage, and hardware-cost constraints across LoRA and FullFT experiments.

  • Dataset statistics: The four datasets span an order of magnitude in token load, setting the achievable-loss floor for the size analysis.Token counts are measured per example and differ across model families because tokenizers differ.
  • Coverage: 970 of 1,008 planned dense cells are eligible, including 682 LoRA and 288 FullFT cells across Qwen3 and Llama.All 38 exclusions are LoRA runs at learning rate 3×10^-3, concentrated on larger models.
  • Coverage: The MoE sweep adds 112 cells, of which 111 are eligible; the sole exclusion is a LoRA run at 3×10^-3 with batch 8.The Qwen3-Next-80B blind test contributes eight eligible single-point cells.
  • Cost context: At one epoch, projected wall-clock is nearly batch-invariant, so batch selection should prioritize loss and memory rather than epoch-time differences.The figure reports LoRA epoch projections at a 5,000-example foundation budget.
  • Cost context: FullFT achieves higher model FLOP/s than LoRA by about the 6N/4N ratio, while both remain far below peak utilization.The comparison averages across datasets and batch sizes and reports spread across cells.
  • Cost context: Peak memory limits batch growth, with FullFT uniformly heavier than LoRA and memory flattening at batch 64 through gradient accumulation.The memory curves are faceted by model and training method.
  • Dataset statistics: Table 3 reports per-dataset token composition and eligible-cell counts by model family.

Appendix C. The iterative-SFT testbed

The testbed uses customer-validated iterative SFT to create internally consistent targets while preserving an explicit boundary: the downstream judge is not independent of training.

  • Construction: Each dataset comes from a production task with a customer-built evaluation, and examples are generated by drafting, grading, feedback, and revision until they pass.The pipeline uses the task evaluator to produce accepted training outputs rather than relying on hand-written examples.
  • Construction: Because every output is accepted by the same evaluator, the supervised targets are internally consistent and changes in validation NLL more directly reflect the lever under study.This reduces label noise and contradictory gold answers that would complicate absolute-loss interpretation.
  • Readout: The customer-validated evaluation serves as the downstream judge because the data was constructed to satisfy it.The setup makes loss and judge comparisons well-posed for the one-lever-at-a-time sweep.
  • Limitation: The judge is not independent: it measures fit to the task as specified rather than serving as an external oracle.The paper uses cross-criterion IFEval to check whether task fit comes with erosion of general capability.

Appendix D. LoRA retention by model

LoRA retains most of FullFT’s improvement with far fewer trainable parameters, while the appendix provides the underlying model, dataset, learning-rate, and size-trend analyses.

  • LoRA retention: The retention breakdown reports, for each model, the fraction of FullFT improvement recovered by LoRA and LoRA’s trainable-parameter percentage.
  • LoRA retention: LoRA recovers a median 98% of FullFT’s gain at only 3–13% of the parameters.Retention is measured over matched dataset-and-batch cells relative to the pretrained baseline.
  • Learning-rate analysis: The learning-rate response curves compare validation NLL across global batches, with LoRA and FullFT shown in separate grids and recommended defaults marked.The six Qwen sizes are split into small and large panels for readability.
  • Size trends: Saturating size-trend fits report L∞, α, bootstrap intervals, exponent refits, and RMSE for each dataset at fixed defaults.Each fit uses six Qwen3 sizes and total parameters measured in billions.
  • Schedule analysis: At 10^-4, constant learning rate preserves a slightly larger step size; at 10^-3, cosine is better calibrated; at 3×10^-3, constant schedules often drift upward late.These comparisons use best-checkpoint loss and end-of-training rebound.

Appendix H. Estimating the Hessian trace: a first-order identity

Appendix H replaces direct Hessian-trace computation with a first-order estimator based on expected squared gradients, while clarifying its saturation-regime limitation and LoRA subspace scope.

  • Motivation and identity: The method estimates converged-loss flatness alongside validation loss using the Hessian trace, which equals the Fisher trace at a well-fit minimiser.The trace is computed as the expected squared gradient norm of a sampled-token log-probability, avoiding second derivatives.
  • Derivation: The derivation splits the log-probability Hessian into a probability-curvature term and an outer-product term, then removes the former by averaging over model-sampled tokens.The cancellation relies on normalised predicted probabilities and the saturation assumption.
  • Estimator and implementation: The resulting estimator uses one backward pass per sample and squared gradient norms, rather than forming the d × d Hessian or using second derivatives.The implementation samples 64 held-out test-split probes per trace and accumulates gradients over trainable parameters.
  • Estimator and implementation: For LoRA, the reported curvature is flatness within the adapter subspace rather than the full parameter space.The fine-tuned grid was re-measured with the same 64-sample test-split setting used for the base-model anchor.
  • Caveat: The estimator is exactly a Hessian trace only at a fully fit minimiser in the saturation regime; otherwise it is a Fisher-trace flatness proxy read alongside loss.This scope boundary is stated explicitly in the regime caveat.

Appendix I. AdamW and Muon: update rules and implementation

Appendix I presents AdamW and Muon in a shared steepest-descent framework, then specifies Muon’s matrix-sign update, approximation, and implementation choices.

  • Update rules: AdamW rescales each scalar parameter independently with a diagonal second-moment preconditioner, whereas Muon applies momentum SGD followed by matrix orthogonalisation to 2D weights.Muon’s update is W_t = W_t−1 − ηs O_t, with O_t = msign(B_t).
  • Matrix sign: The matrix sign sets every singular value to 1, equivalently producing the nearest semi-orthogonal matrix and the steepest direction under a bounded operator norm.These are three equivalent characterisations of msign.
  • Implementation: Muon approximates the matrix sign with Newton–Schulz matrix multiplies that independently drive singular values toward 1 within the iteration’s basin.The fixed odd-polynomial map is stable on tensor cores after spectral normalisation.
  • Shared geometry: AdamW and Muon are framed as steepest descent under different norms: elementwise max-norm for AdamW and spectral norm for Muon.For a matrix layer, operator-norm control bounds the worst-case output change.
  • Related optimisers: The appendix relates Muon to Shampoo and notes that their single-gradient limit yields the same orthogonalised step, while accumulated second-moment information differs.SOAP is described as an intermediate approach.
  • Implementation choices: Muon’s behaviour depends on tensor-parallel orthogonalisation, QKV splitting, scale matching, and a hybrid rule that sends vectors, embeddings, and output heads to Adam.Under LoRA, Muon orthogonalises 2D adapter matrices, matching the curvature scope used in Appendix H.
Loading 2609.01244v1…