Adaptive-Depth Is a Finishing Stage
Organizations: University of Pennsylvania · University of Maryland
Abstract
Looped transformers reuse the same layer multiple times, allowing inference compute to be adjusted without changing the model. Recent adaptive-depth methods train the model across multiple loop depths from the beginning of training. We introduce fixed-first training, which varies model depth only during a short final phase of training. At 200M parameters, training a model of fixed-depth 2 for the first 67% of the run reduces modeled training compute by 45% while keeping mean negative log-likelihood (NLL) across inference depths within 0.5% of training across depths throughout. At 1B parameters, training initially at depth 4 reduces modeled compute by 31% to 37% and wall-clock time by 32% to 39%, while keeping final NLL within 0.8% and downstream differences within the observed seed variation. These results show that training over the full range of depths can be reserved for a short finishing phase with only small changes in performance.
Figures & tables
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
| Part | Token sequence |
|---|---|
| Person’s name | N17 N204 N6 |
| Birth city | V3 V11 |
| Employer | V6 V10 |
| Stage | What the model sees |
|---|---|
| Training | N17 N204 N6 [employer] V6 V10 [birth city] V3 V11 |
| Evaluation | N17 N204 N6 [birth city] ___ ___ [employer] V6 V10 |
| Correct response | V3 V11 |
| 200M schedule | Terminal NLL | Mean serving NLL | Modeled FLOPs | Wall-clock reduction |
|---|---|---|---|---|
| Uniform | 0.67297 0.00058 | 0.68582 0.00077 | 100% | – |
| Fixed 67% | 0.67424 0.00027 | 0.68643 0.00028 | 68.7% | 26.2% |
| Uniform, matched | 0.71936 0.00066 | 0.72089 0.00072 | 68.7% | – |
| Fixed 80% | 0.67677 0.00037 | 0.68941 0.00087 | 62.4% | 31.6% |
| Depth-averaged task | Uniform | Fixed 67% | Fixed 80% |
|---|---|---|---|
| HellaSwag | 50.72 0.33 | 50.56 0.19 | 50.27 0.24 |
| PIQA | 70.48 0.42 | 70.52 0.55 | 70.86 0.53 |
| ARC-Easy | 56.11 1.46 | 56.76 1.17 | 57.43 1.12 |
Efficient Language Model Training
Efficient language model training studies how to reduce the data, compute, memory, time, energy, and cost needed to build capable language models.
See all 162 papers