Accordion: Elastic-Depth Inference via Fixed-Prefix Layer Looping

Summary (Overview)

  • Accordion is a mixture-of-experts (MoE) Transformer that supports five inference budgets (Low, Medium, High, XHigh, Ultra) through fixed-prefix layer looping, using a single set of shared weights across all budgets.
  • The architecture repeats a fixed prefix of layers 7–18 within a middle window, with zero-initialized residual gates controlling second-pass computations. The minimum budget (k=0) retains the unlooped model, avoiding the "one-loop tax" problem.
  • After 3.6T tokens of pretraining, increasing inference depth improves the five-category macro-average across 37 tasks by 3.01 percentage points, with the middle budget (k=6) capturing 89% of this gain using half as many additional layer executions.
  • The training pipeline extends elastic-depth inference through supervised fine-tuning (SFT) and reinforcement learning (RL), integrating three domain experts (instruction following, mathematics, code) via parameter-group merging and budget-matched multi-teacher on-policy distillation (MOPD).
  • The final model (Accordion-RL) surpasses each domain expert on its corresponding task-and-budget average and outperforms four public reference models on IFEval at every inference budget.

Introduction and Theoretical Foundation

Background and Motivation

Large language models serve applications with diverse inference budgets, but training and maintaining separate models for each budget incurs substantial costs. Elastic inference aims to address this by supporting multiple compute choices within a single model. Existing approaches include:

  • Nested architectures and subnetworks (Devvrit et al., 2023; Cai et al., 2024)
  • Reduced execution depth (Fan et al., 2019; Schuster et al., 2022)
  • Input-dependent computation allocation (Raposo et al., 2024; Bae et al., 2025)

Looped Transformers

Looped Transformers repeatedly apply shared parameters to increase effective depth. Key challenges include:

  • The "one-loop tax" problem (Ladhak et al., 2026): jointly optimizing multiple inference depths can compromise shallow-path performance
  • Maintaining elastic inference through the full training pipeline (pretraining, SFT, RL)

Key Theoretical Foundation

Accordion builds on layer looping from Loopie (Gao et al., 2026), which repeats selected Transformer layers consecutively rather than repeating an entire model. The architecture uses a fixed-prefix layer selection strategy, where each budget repeats a prefix of a middle window of layers.


Methodology

Architecture

Consider a pre-norm MoE transformer with LL physical layers and hidden dimension dd. Let W={w1,…,wm}\mathcal{W} = \{w_1, \ldots, w_m\} be an ordered, contiguous middle window, and K⊆{0,…,m}\mathcal{K} \subseteq \{0, \ldots, m\} the budget menu. Budget kk repeats the prefix Sk={w1,…,wk}S_k = \{w_1, \ldots, w_k\}, with S0=∅S_0 = \varnothing. Its path PkP_k executes layer ℓ\ell exactly:

rℓ(k)=1+1{ℓ∈Sk}(1)r_{\ell}(k) = 1 + \mathbf{1}\{\ell \in \mathcal{S}_k\}\tag{1}

times before moving to layer ℓ+1\ell + 1.

Configuration used: L=24L = 24, W={7,…,18}\mathcal{W} = \{7, \ldots, 18\}, K={0,3,6,9,12}\mathcal{K} = \{0, 3, 6, 9, 12\}

Training Objective

At each training micro-batch, a budget k∼Uniform(K)k \sim \text{Uniform}(\mathcal{K}) is sampled, and only path PkP_k is executed. The next-token prediction loss is:

LNTP=−1N∑x∈B∑tlog⁡pθ(xt∣x<t;Pk)(2)\mathcal{L}_{\mathrm{NTP}} = -\frac{1}{N}\sum_{x \in \mathcal{B}}\sum_{t}\log p_{\theta}(x_t \mid x_{<t}; P_k)\tag{2}

The full training objective includes a sequence-level MoE balancing term:

J(θ)=EB,k[LNTP+λseqLseq](3)\mathcal{J}(\theta) = \mathbb{E}_{\mathcal{B}, k}[\mathcal{L}_{\mathrm{NTP}} + \lambda_{\mathrm{seq}}\mathcal{L}_{\mathrm{seq}}]\tag{3}

with λseq=10−4\lambda_{\mathrm{seq}} = 10^{-4}.

Layer Looping Mechanism

For layer input hh, the first pass is:

u=h+Aℓ(NℓA(h)),v=u+Mℓ(NℓM(u);bℓ(1)).(4)\begin{array}{l} u = h + A_{\ell}(\mathrm{N}_{\ell}^{A}(h)), \\ v = u + M_{\ell}(\mathrm{N}_{\ell}^{M}(u); b_{\ell}^{(1)}). \end{array}\tag{4}

If ℓ∉Sk\ell \notin S_k, the layer returns vv. Otherwise, it immediately applies a second pass:

u2=v+tanh⁡(gℓA)⊙Aℓ(NℓA(v)),v2=u2+tanh⁡(gℓM)⊙Mℓ(NℓM(u2);bℓ(2)).(5)\begin{array}{l} u_2 = v + \tanh(g_{\ell}^{A}) \odot A_{\ell}(\mathrm{N}_{\ell}^{A}(v)), \\ v_2 = u_2 + \tanh(g_{\ell}^{M}) \odot M_{\ell}(\mathrm{N}_{\ell}^{M}(u_2); b_{\ell}^{(2)}). \end{array}\tag{5}

The per-channel gates gℓA,gℓM∈Rdg_{\ell}^{A}, g_{\ell}^{M} \in \mathbb{R}^d are zero-initialized and excluded from weight decay, making the second pass an identity map at initialization.

Computational Cost

For equal-shaped layers, FLOPs are approximately:

C(k)≃Cfixed+(L+k)Clayer(6)C(k) \simeq C_{\mathrm{fixed}} + (L + k) C_{\mathrm{layer}}\tag{6}

KV cache storage scales as (L+k)/L(L+k)/L relative to k=0k=0: 1.25×1.25\times at k=6k=6 and 1.5×1.5\times at k=12k=12.

Design Studies (1.61B-parameter MoE)

Four key design choices were evaluated:

  1. Layer selection: Fixed prefix (7–18) outperforms fixed center, fixed suffix, and randomized strategies at all looped budgets
  2. Number of budgets: Five budgets outperform three or four at the same expected training computation
  3. Router sharing: Shared router achieves lowest mean loss (improvement of 0.0013 nats over separate routing)
  4. Budget-conditioned parameters: Budget-independent gates and norms achieve the lowest loss at every budget

Training Pipeline

StageTokens (B)CumulativeContextLearning Rate
General3,0003.00T8K5×10−45 \times 10^{-4}, constant
Mid-training1503.15T8K5×10−45 \times 10^{-4}, constant
Annealing4503.60T8K5×10−4→5×10−55 \times 10^{-4} \rightarrow 5 \times 10^{-5}
Context extension503.65T32K5×10−5→1.5×10−55 \times 10^{-5} \rightarrow 1.5 \times 10^{-5}

Model Configuration (Accordion-6B-A0.5B)

SettingValue
Layers / hidden size / Q / KV / head dim24 / 1536 / 12 / 4 / 128
Routed / shared / active experts128 / 1 / 8
Total / active parameters6.03B / 0.54B
Loopable window / budgets7–18 / {0, 3, 6, 9, 12}
Peak LR / final LR5×10−45 \times 10^{-4} / 10% of peak
AdamW β / weight decay(0.9, 0.95) / 0.1

Reinforcement Learning Pipeline

Three domain experts are trained using GRPO/DAPO-style training, then combined via:

  1. Parameter-group merging (task arithmetic):
θqmerge=θSFT,q+∑jλj,q(θj,q−θSFT,q)\theta_{q}^{\mathrm{merge}} = \theta_{\mathrm{SFT}, q} + \sum_{j} \lambda_{j, q}(\theta_{j, q} - \theta_{\mathrm{SFT}, q})
  1. Budget-matched multi-teacher on-policy distillation (MOPD): Each prompt's teacher is chosen by domain and evaluated at the same loop budget as the student

Empirical Validation / Results

Design Study Results

Table 1: Comparison of training-budget sets (validation loss in nats)

ModelTraining budgetsk=0k=6k=12Mean
Five-budget model{0,3,6,9,12}2.57932.57032.56962.5744
Three-budget model{0,6,12}2.58102.57292.57212.5766
Four-budget model{0,4,8,12}2.5814-2.57192.5767

Table 3: Comparison of budget-conditioned parameters (validation loss in nats)

Configurationk=0k=3k=6k=9k=12Mean
Budget-independent gates and norms2.57932.57312.57032.56982.56962.5724
Budget-conditioned gates2.58142.57612.57312.57232.57212.5750
Budget-conditioned norm scales2.58302.57692.57352.57242.57222.5756

Specialist Comparison

Under equal training tokens (80B), Accordion achieves validation losses within 0.0030 nats of budget-specific specialists at all three evaluated budgets (k=0, 6, 12). Under equal training FLOPs, it improves on the FLOPs-matched specialist at k=12 by 0.0068 nats.

Base Model Results (37 Tasks)

Increasing budget from k=0 to k=12:

  • Raises 37-task aggregate from 47.05 to 50.06 (+3.01 percentage points)
  • Improves 36 of 37 task scores
  • Middle budget (k=6) captures 89% of endpoint gain with ~21–24% increase in time per output token

Category averages (Figure 7):

  • Largest gains in code, reasoning, and mathematics
  • Ranks second in knowledge/QA and common sense at k=12
  • Ranks first in code at k=9

Selected Benchmark Results (Accordion-Base vs. References)

Benchmarkk=0k=6k=12Qwen3-0.6BQwen3-1.7BOLMoE-0125Ouro-1.4B
MMLU58.260.960.952.562.755.267.5
BBH54.860.060.840.853.332.670.8
GSM8K53.159.461.359.271.351.178.6
MATH (5,000)39.842.242.428.538.916.922.0
HumanEval48.851.253.128.750.611.635.37
MBPP46.650.251.839.055.621.058.0

Post-Training Results

Accordion-Instruct (SFT):

  • Increasing budget improves 9 of 10 task scores
  • Exceeds Qwen3-1.7B on HumanEval+ at k=12 (by 2.44 points), MBPP+ at k=6, and AIME 2025 at k=9

Accordion-RL (after RL and expert consolidation):

Benchmarkk=0k=6k=12Qwen3-0.6BQwen3-1.7B
IFEval87.9989.8389.8362.8571.35
MMLU-Redux70.9874.1374.7757.6974.84
GSM8K84.1588.6389.1679.4590.14
MATH50083.8088.6089.8072.8089.20
AIME 202434.5840.8345.2112.0844.58
HumanEval+79.8886.5984.1546.3479.27

Key findings:

  • Exceeds all four public references on IFEval at every budget (margin ≥16.64 points over strongest reference)
  • All 15 task–budget comparisons on IFEval, IFBench, and MMLU-Redux show higher scores after RL vs. SFT
  • Exceeds each domain expert on its corresponding task-and-budget average (paired 95% intervals excluding zero)

Long-Context Results (RULER)

  • At 4K context: +0.72 points from k=0 to k=12
  • At 32K context: +5.47 points from k=0 to k=12
  • At 16K, k=3 captures nearly all endpoint gain; at 32K, larger budgets unlock further gains through k=9

Mechanistic Analysis (Second Pass)

  • Layer 9 retains only 1.52 of 8 routed experts on average between passes; most tokens share at most one expert
  • Second-pass MoE writes are comparable in size to first-pass writes but nearly orthogonal to them
  • Residual-stream cosines stay above 0.967 across the repeated window
  • Removing adjacent repetitions together costs 15–33% of the summed individual effects (Hydra effect), while spaced triplets cost 96–103%

Theoretical and Practical Implications

Theoretical Implications

  1. Fixed-prefix layer looping provides an effective mechanism for elastic-depth inference without the "one-loop tax," preserving unlooped path quality while enabling depth-dependent gains.

  2. Uniform budget sampling during pretraining is sufficient—no cross-budget consistency objectives are needed, simplifying the training objective.

  3. Shared parameters support distinct computations: The second pass performs input-dependent expert reselection and writes distinct residual updates, demonstrating that repeated layers do not simply duplicate first-pass computation.

  4. Non-additive effects of repetition: The Hydra effect (joint removal costing less than sum of individual removals) for adjacent layers shows that depth benefits must be assessed on complete budget paths, not by summing individual layer contributions.

Practical Implications

  1. One model, five inference budgets: Accordion provides flexible quality–cost trade-offs for diverse application needs, eliminating the need to train and maintain separate models.

  2. Full-pipeline elastic inference: The approach extends through SFT and RL, making elastic depth practical for production deployment.

  3. Public-data recipe at scale: The complete training pipeline uses publicly available datasets, enabling reproducibility and further research.

  4. Task-specific operating points: Different tasks peak at different budgets (e.g., HumanEval at k=9, MBPP+ at k=12), allowing users to select the optimal budget per workload.


Conclusion

Accordion makes inference depth a selectable property of a single model through fixed-prefix layer looping. Key achievements:

  • Five trained execution paths sharing one set of weights, with the minimum budget retaining the unlooped model
  • Depth gains preserved through pretraining, SFT, and RL, with the base model obtaining most gains at the middle budget
  • Budget-matched multi-teacher distillation combines three domain experts into a student exceeding each expert on its domain's task-and-budget average
  • Strong instruction-following outperforming four public references on IFEval at every budget
  • Mechanistic insights showing repeated layers reselect experts and write distinct updates, with prediction effects depending on neighboring repetitions

Future directions may include exploring alternative budget menus, extending to even larger scales, and investigating whether learned (rather than fixed) path selection can further improve the quality–cost frontier.

Related papers