Hyperparameter Search:
SFT & DPO for Grade-School Math

Temperature scales the model's raw output probabilities (logits) before sampling. A value of 1.0 leaves probabilities unchanged; below 1.0 sharpens the distribution so the top token dominates; above 1.0 flattens it so rare tokens get more chances. For Pass@5 (where you generate five independent answers and check if any is correct), a little randomness is your friend: it lets the model explore different solution paths across the five samples.

Range rationale: [0.6, 1.4] is what we found to be a practical bracket for sampling diversity without sacrificing coherence. The 0.6 floor ensures at least some variation between samples; the 1.4 ceiling keeps outputs from degenerating into noise.

Repetition penalty down-weights any token that has already appeared in the generated text, making it less likely to be chosen again. A value of 1.0 applies no penalty at all; values above 1.0 increase the suppression proportionally. For math reasoning, this primarily prevents the model from getting stuck in looping phrases like "so we add 4, so we add 4, so we add 4…"

Range rationale: [1.0, 1.3] is a narrow, conservative window that we think is appropriate here. Math solutions inherently repeat structural keywords ("so", "therefore", "the answer"); the goal is to catch genuine infinite loops without penalising those legitimate repetitions.

Learning rate is the "step size" of the optimizer: after computing how the loss changes with respect to each weight, the model moves lr × gradient in the corrective direction. We search this in log-space because the right order of magnitude matters far more than a precise decimal value. For instance, 1e-4 and 2e-4 behave similarly, while 1e-4 and 1e-3 can behave very differently.

Range rationale: [5e-5, 8e-4] is what we considered a practical window for fine-tuning a small transformer on this task. Pre-trained weights are delicate; the goal is to nudge them, not overwrite them.

Each epoch is one complete pass through the training dataset. More epochs mean more opportunities to learn the target format, but with diminishing returns and increasing risk of memorisation.

Range rationale: [1, 5] reflects both the small size of the GSM-8K training split we use and our Colab T4 (16 GB) session budget. With a compact dataset, the model tends to saturate quickly; beyond epoch 5, you are mostly memorising rather than generalising, and the time cost becomes hard to justify given the session limits.

Batch size determines how many training examples the model processes before computing a gradient update. Larger batches produce smoother, more reliable gradient estimates; smaller batches are noisier but update more frequently, and can act as an implicit regulariser.

Range rationale: {4, 8, 16} covers what we think is a practical range given the strict 16 GB GPU memory limit of a Colab T4 and the 512-token sequence length fixed for SFT. All three are powers of two for memory alignment efficiency.

Weight decay (L2 regularisation) adds a small penalty proportional to the magnitude of each weight to the loss, nudging the optimiser to prefer smaller, more diffuse parameter values. This discourages the model from leaning too hard on any single feature, which can help generalisation.

Range rationale: [0.0, 0.2] felt like a practical window to us for this task. 0.0 is a perfectly valid choice (no regularisation at all), and we found values beyond 0.2 tend to over-constrain the model without clear benefit.

DPO uses a learning rate orders of magnitude smaller than SFT, typically somewhere in the 1e-7 to 1e-5 range rather than 1e-4 to 1e-3. This is intentional: DPO is meant to make subtle preference-driven adjustments on top of an already capable model, not rewrite it from scratch.

Range rationale: [1e-7, 1e-5] spans three orders of magnitude specifically to explore the boundary between "barely moving" and "moving enough to matter". Because any individual step is so small at the lower end, the model needs more passes to accumulate a meaningful signal, which is why the DPO epoch range extends to 10 rather than 5.

The DPO epoch range is intentionally wider than SFT's [1, 5] to compensate for the much lower learning rate. When each individual update is tiny (e.g. at lr=1e-7), the model needs many more passes through the preference data before the signal accumulates into a meaningful shift in behaviour.

Range rationale: [1, 10] is set wide to cover both ends of the LR spectrum. At a higher LR (1e-5), a single epoch may already be enough; at a lower LR (1e-7), you need closer to 10 epochs for meaningful change. Keeping the upper bound at 10 also kept total DPO training time within our Colab T4 (16 GB) session budget.

Beta (β) is the KL-regularisation coefficient in the DPO objective. It controls how strongly the training penalises the policy model for drifting too far from the reference model (the SFT checkpoint). A higher β keeps the new policy close to the reference; a lower β allows larger deviations in pursuit of the preference signal.

Range rationale: [0.05, 0.4] is what we considered a practical range where the KL constraint is meaningful but not overwhelming. Values around 0.1–0.2 appear frequently in DPO literature for small-scale experiments, which informed our choice of upper bound.

DPO requires both the live policy model and the frozen reference (SFT) model to be in memory simultaneously during training, roughly doubling the GPU footprint compared to SFT. This is why the batch size options here are smaller than in SFT; 16 would exceed the strict 16 GB VRAM limit of the Colab T4 under the dual-model load.

Range rationale: {4, 8} are the options we found fit comfortably within the Colab T4's 16 GB memory budget given the dual-model overhead of DPO and the 512-token sequence length.

Same L2 regularisation mechanism as in SFT. During DPO, with such a small learning rate, weight decay plays a more prominent relative role; even a modest value of 0.05 can meaningfully constrain the magnitude of the preference-driven updates.

Range rationale: [0.0, 0.2] mirrors the SFT range. 0.0 is a perfectly reasonable default when the learning rate alone provides sufficient implicit regularisation, and we saw little benefit from going beyond 0.2.