Adaptive-Depth Is a Finishing Stage

arXiv:2610.14252 · cs.LG · Submitted 2026-10-09 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom.

Jane: And I'm Jane, here with Lu, Meng and our in-house large language model, Lalam.

Tom: Today's paper: "Adaptive-Depth Is a Finishing Stage".

Jane: The gist The fixed-first training method allows adaptive-depth looped transformers to adjust model depth only during a short final phase of training,

Tom: First, who's behind it and why it matters.

Paper summary: Tom: We're looking at the paper "Adaptive-Depth Is a Finishing Stage," and it lays out the core idea of this fixed-first training strategy for adaptive-depth looped transformers.

Lu: The authors are proposing that instead of changing the model depth all through training, you only vary it during a short final phase of training after an initial fixed period.

Jane: It suggests that this two-stage approach lets you get fast learning and specialization at the beginning, and then quickly fine-tune those depths at the very end.

Conclusion: Tom: We're wrapping up on "Adaptive-Depth Is a Finishing Stage." The authors are Zeyuan Allen-Zhu and Yuanzhi Li.

Jane: It’s about using this fixed-first training method to save compute while still getting the benefits of adaptive depth in those looped transformers.

Lu: What they're really proposing is that you only adjust the model depth during a very short final phase of training instead of having it change all the time at the start.

Meng: So it’s like building something fast and stable first, and then just tweaking the final details later to get that deep model we want.

Tom: They show this works across different sizes, like at two hundred million parameters where they saw a forty-five percent reduction in modeled compute.

Jane: And for a one billion parameter model, they saw savings of thirty-one to thirty-seven percent during training.

Lu: The key finding is that the training over the full range of depths can be reserved for that short finishing phase with only small changes in performance.

Tom: That means you get a lot more compute budget available for other things when you're working on these deep models.

Jane: Overall, this paper suggests a way to manage the complexity of adaptive depth while keeping the training process efficient and focused on getting useful recall sooner.

Lu: It's about finding that sweet spot between making the model flexible enough and keeping it computationally affordable across all those different depths.

Meng: This means we can build these advanced transformers without needing a massive compute budget right away to get started.

Andrei Piterbarg, Sean McLeish, Dhruv Pai, Benjamin Keigwin

University of Pennsylvania · University of Maryland

cs.LG

Submitted: 2026-10-09

Updated: 2026-10-09

Comments: Accepted for poster presentation at the AXIOM Workshop (Foundations of Efficient Deep Learning) at NeurIPS 2026. 9 pages, 8 figures

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 89/100

The gist: The gist The fixed-first training method allows adaptive-depth looped transformers to adjust model depth only during a short final phase of training, significantly reducing compute while maintaining

Key concepts

Fixed-First Training: This technique trains the transformer model at a fixed depth for the majority of the training run. Depth is only adjusted during a brief final phase. This method aims to achieve quick initial learning while ensuring that the model retains good performance across different required serving depths, leading to substantial compute savings.
Adaptive-Depth Looped Transformers: These are specialized transformer architectures designed to handle varying levels of depth efficiently. They allow the model's structure (its depth) to change adaptively during training. The fixed-first method controls this adaptation by limiting when and how much the depth is allowed to change.
Serving Depth NLL: This metric measures how well the model performs when deployed at a specific required 'serving depth.' The paper shows that fixed-first training keeps this value very close to uniform training, meaning it maintains high performance even when the model needs to operate at different structural depths.
Finishing Stage Length Diagnostic: This diagnostic investigates how long the final adjustment phase should be. It suggests that while deep updates are useful for flexibility, their advantage diminishes quickly after a certain number of uniform updates, indicating that a short finishing stage is sufficient for achieving the desired performance gains.

Summary

The gist The fixed-first training method allows adaptive-depth looped transformers to adjust model depth only during a short final phase of training, significantly reducing compute while maintaining performance across serving depths.

Motivation for Fixed-First Training

The motivation for fixed-first training stems from a controlled factual-recall experiment that demonstrates the two stages of learning: fixed-depth training learns quickly but specializes near its training depth, while uniform training retains recall over a broader range of depths but learns more slowly. This is tested using a synthetic memorization task where models are evaluated on the exact answer across five fixed templates. The comparison shows that fixed-first training achieves useful recall sooner and retains more of that recall beyond the depth range used during training

Key Findings on Model Size

The benefits of fixed-first training are observed across different model scales, with substantial compute savings reported At 200M parameters, training at fixed depth 2 for the first 67% of the run reduces modeled compute by 45% For a 1B parameter model, fixed-first training reduces modeled training compute by 31% to 37% 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

Performance Across Serving Depths

The paper compares fixed-first training against uniform-depth training and progressive Poisson-lognormal curricula At 200M, using K0 = 2 for 67% of training saves 45% of modeled compute while keeping mean serving-depth NLL within 0.5% of uniform training At 1B, the schedules show that the K0 = 4, 67% schedule nearly matches full uniform training in mean serving-depth NLL (0.68643 versus 0.68582) with 31.3% less modeled compute

Diagnostic on Finishing Stage Length

The research investigates the short duration of the finishing phase by comparing single updates at K=4 and K=16 The diagnostic suggests that deep updates remain important for flexibility, but their relative advantage becomes small soon after the uniform phase begins This occurs because the depth-averaged advantage of the K=16 update is no longer distinguishable from zero after 16 uniform updates

Language Model Results

When applied to language modeling tasks, fixed-first training yields significant efficiency gains For the 1B model, the 67% and 80% schedules cut training FLOPs by 30.8% and 36.9% The terminal NLL remains within a specified margin of error relative to uniform training Overall, fixed-first training cuts 30% to 40% of modeled training compute and 26% to 39% of wall-clock time while keeping NLL within 1%

Conclusion on Efficiency

The study concludes that fixed-first training is an effective strategy for adaptive-depth looped transformers, demonstrating substantial reductions in both modeled compute and wall-clock time The method successfully balances the need for depth specialization with the efficiency of a shorter finishing phase

References

Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 3.1, knowledge storage and extraction. arXiv preprint arXiv:2309.14316, 2023 <ref:2610.14252#pg6>

Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 7432–7439, 2020 <ref:2610.14252#pg6>

Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457, 2018 <ref:2610.14252#pg6>

Jonas Geiping, Sean McLeish, Neel Jain, John Kirchenbauer, Siddharth Singh, Brian R Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, and Tom Goldstein. Scaling up test-time compute with latent reasoning: A recurrent depth approach. arXiv preprint arXiv:2502.05171, 2025 <ref:2610.14252#pg6>

Amr Hegazy, Amr Alanwar, and Mostafa Elhoushi. Gated recurrent transformers: Expressive depth through recurrent modulation in transformers. arXiv preprint arXiv:2608.15062, 2026 doi: 10.48550/arXiv.2608.15062 <ref:2610.14252#pg8>

Ahmadreza Jeddi, Marco Ciccone, and Babak Taati. Loopformer: Elastic-depth looped transformers for latent reasoning via shortcut modulation. arXiv preprint arXiv:2602.11451, 2026 <ref:2610.14252#pg6>

Sean McLeish, Ang Li, John Kirchenbauer, Dayal Singh Kalra, Brian R Bartoldson, Bhavya Kailkhura, Avi Schwarzschild, Jonas Geiping, Tom Goldstein, and Micah Goldblum. Teaching pretrained language models to think deeper with retrofitted recurrence. arXiv preprint arXiv:2511.07384 <ref:2610.14252#pg6>

Nikunj Saunshi, Nishanth Dikkala, Zhiyuan Li, Sanjiv Kumar, and Sashank J Reddi. Reasoning with latent thoughts: On the power of looped transformers. arXiv preprint arXiv:2502.17416 <ref:2610.14252#pg6>

Avi Schwarzschild, Eitan Borgnia, Arjun Gupta, Furong Huang, Uzi Vishkin, Micah Goldblum, and Tom Goldstein. Can you learn an algorithm? generalizing from easy to hard problems with recurrent networks. Advances in Neural Information Processing Systems, 34:6695–6706, 2021 <ref:2610.14252#pg9>

Yudong Wang, Zixuan Fu, Jie Cai, Peijun Tang, Hongya Lyu, Yewei Fang, Zhi Zheng, Jie Zhou, Guoyang Zeng, Chaojun Xiao et al. Ultra-fineweb: Efficient data filtering and verification for high-quality llm training data. arXiv preprint arXiv:2505.05427 <ref:2610.14252#pg7>

Yudong Wang, Zixuan Fu, Hengyu Zhao, Chen Zhao, Chuyue Zhou, Xinle Lin, Hongya Lyu, Shuaikang Xue et al. Data science and technology towards agi part i: Tiered data management. arXiv preprint arXiv:2602.09003 <ref:2610.14252#pg6>

Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? arXiv preprint arXiv:1905.07830 <ref:2610.14252#pg7>

Rui-Jie Zhu, Zixuan Wang, Kai Hua, Tianyu Zhang, Ziniu Li, Haoran Que, Boyi Wei et al. Scaling latent reasoning via looped language models. arXiv preprint arXiv:2510.25741 <ref:2610.14252#pg6>

Yudong Wang, Zixuan Fu, Hengyu Zhao, Chen Zhao, Chuyue Zhou, Xinle Lin, Hongya Lyu et al. Data science and technology towards agi part i: Tiered data management. arXiv preprint arXiv:2602.09003 <ref:2610.14252#pg6>

Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag:

Improvements for AI systems

The system can achieve significant inference compute reduction while maintaining high performance across varied serving depths by implementing a fixed-first training strategy, which train[s] at an initial depth K0 for the first 67% or 80% of training, then switches to uniform K=1,6 training. This reduces modeled compute by 31% to 45% while keeping mean negative log-likelihood (NLL) across inference depths within 0.5% (at 200M parameters).
The AI system can be optimized for faster deployment without sacrificing downstream accuracy by leveraging the observed behavior that deep updates remain important for flexibility, but their relative advantage becomes small soon after the uniform phase begins. This suggests that a shorter finishing phase is sufficient to correct depth specialization once a strong low-depth computation is present, leading to reduced wall-clock time.
The model can be specialized for tasks requiring both fast initial learning and broad generalization by employing a staged training approach where the initial fixed depth learn[s] quickly but specializes near its training depth, followed by a uniform phase to ensure recall over a broader range of depths.

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.