What question did this study set out to answer?

This research investigates how Adafactor's step-size characteristics can stall training in large transformers.

July 12, 20260 citationsOpen Access

When Parameter Scaling Stalls Training: Diagnosing Adafactor Step-Size Semantics in a 336M-Parameter Transformer

Key Points

This research investigates how Adafactor's step-size characteristics can stall training in large transformers.
Pretrained a 336M-parameter transformer on a mixed 891M-token corpus using 16 TPU v6e chips.
Enabled parameter scaling without retuning the global learning rate and recorded update RMS and loss reduction.
Conducted a parallel weight-decay study to compare effects.
Enabling parameter scaling reduced mean global update RMS to 3.05% of control, changing loss reduction from 2.016 to 0.016.
Validation loss was 9.7404 for scaled parameters versus 5.5322 for unscaled control.
A restored run at 4.3 × 10−3 recaptured 96.5% of control update RMS, with a correlation of 0.996 to control loss trajectory.

Abstract

Adafactor can express either an absolute step size or a step size multiplied by eachparameter block’s root-mean-square (RMS) magnitude. Enabling the latter withoutretuning the global learning rate therefore changes the units of the optimizationstep. We study this failure mode while pretraining a 336M-parameter decoder-onlyTransformer on a mixed 891M-token corpus using 16 TPU v6e chips. In a matchedten-step diagnostic, enabling parameter scaling while retaining the absolute-steplearning rate of 10−4 reduced mean global update RMS to 3.05% of the controland changed the loss reduction from 2.016 to 0.016. A two-epoch run consequentlyreached validation loss 9.7404, compared with 5.5322 for the unscaled control. Theinitial global parameter RMS was 0.02343, suggesting the heuristic compensation10−4/0.02343 ≈ 4.27 × 10−3 . A fresh run at 4.3 × 10−3 restored 96.5% of thecontrol update RMS; its ten-step loss trajectory had correlation 0.996 and meanabsolute deviation 0.0459 from the control. A parallel weight-decay study foundno clear benefit over the controlled first ten token-budget epochs. These resultsidentify update suppression, rather than numerical divergence, as the observedfailure mechanism and motivate reporting update-to-parameter scale wheneverAdafactor parameter scaling is changed.

Ask AI

Helpful

Bookmark

View Full Paper