Pith. sign in

REVIEW 3 major objections 56 references

A single learned controller that jointly sets attention mode, experts, and cache bits for each token beats the best independent mix of MoD, MoE, and KV quantization at the same inference budget.

Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

T0 review · grok-4.5

2026-07-11 03:19 UTC pith:KIDMCY2P

load-bearing objection Solid joint-routing formulation and collapse diagnosis; the Pareto/tail numbers are disclaimed as protocol sketches, so the performance claim is not yet evidence. the 3 major comments →

arxiv 2607.06601 v1 pith:KIDMCY2P submitted 2026-07-07 cs.LG cs.AI

TriRoute: Unified Learned Routing for Joint Adaptive Attention, Experts, and KV-Cache Allocation

Andrii Balashov , Olena Ponomarova This is my paper
classification cs.LG cs.AI
keywords conditional computationmixture of expertsmixture of depthsKV cache quantizationadaptive attentionlearned routinglanguage modelsinference efficiency
verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

Language models normally spend the same compute and memory on every token, even though many tokens are predictable function words and a few rare entities carry most of the meaning. Three popular ways to cut that cost—routing to sparse experts, skipping whole layers, and storing the attention cache at low precision—have been developed and tuned separately. This paper argues those three choices are coupled: a rare name may need full attention and high-precision storage even if it can skip the feed-forward layer. It introduces TriRoute, a small shared controller that, for every token at every layer, jointly picks an attention mode, a set of experts (including a null skip), and a cache bit-width, trained under one global compute-and-memory budget. On decoder-only models from 160M to 1.3B parameters the joint controller matches dense quality at roughly half the inference cost and, unlike the best independent combination, largely preserves accuracy on rare entities, code, and arithmetic.

Core claim

A shared learned controller that jointly sets, per token and per layer, attention resolution (skip/local/full), FFN expert selection (including a null expert that recovers depth-skipping), and KV-cache bit-width under a single Lagrangian budget produces better quality–cost trade-offs than the best independently tuned combination of Mixture-of-Depths, Mixture-of-Experts, and KV quantization, while better protecting tail-case performance that pure average-perplexity optimization erodes.

What carries the argument

TriRoute: a lightweight shared trunk with three heads that emits a coordinated per-token, per-layer policy over attention mode, top-k experts (null expert recovers MoD), and cache bit-width; trained with heterogeneous straight-through Gumbel estimators, per-axis whitening and entropy floors that stop a cross-axis collapse cascade, and an online Lagrangian that turns average FLOPs and memory into one controllable knob.

Load-bearing premise

The reported quality and cost numbers come from completed end-to-end training runs at the stated scales, not from protocol sketches or illustrative curves.

What would settle it

Train the 1.3B TriRoute model and the independently tuned MoD+MoE+KV-quant combination from scratch under identical compute-optimal token budgets, measure Pile validation perplexity and rare-entity/code/math bucket perplexities at the matched (0.55, 0.40) FLOPs/memory point, and check whether TriRoute still lies strictly below the independent combination on both average and tail metrics.

Watch this falsifier — get emailed when new claim-graph text bears on it.

If this is right

  • At a fixed inference budget, joint routing recovers more of the dense model’s quality than separately tuned sparsity, depth, and quantization knobs.
  • Tail robustness on rare entities, code, and arithmetic need not be sacrificed when compute is cut, if the controller can keep full attention and high-precision cache on those tokens while economizing elsewhere.
  • A single budget scalar can replace three hand-tuned sparsity and precision hyperparameters for practitioners.
  • The learned policy is interpretable: full attention and high-precision cache concentrate on sentence starts, rare subwords, and named entities, while FFN compute decouples and focuses on content and numbers.
  • The same controller template can add further axes, such as whether to keep a token’s KV entry at all.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • If the coordination benefit holds at larger scales, production serving stacks could replace three independent sparsity systems with one router, simplifying deployment and budget tuning.
  • The observed “attend early, compute late” schedule suggests hard-coded layer schedules might be replaceable by learned per-token policies without losing early-context benefits.
  • Because attention and bits couple tightly while FFN routing is more independent, systems work on mixed-precision caches and ragged attention may unlock more wall-clock gain than further expert-routing improvements alone.
  • The collapse-cascade fix (per-axis whitening plus entropy floor) may transfer to any multi-axis discrete routing problem where a shared feature trunk can starve one head of variance.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit.

Referee Report

3 major / 0 minor

Summary. The paper proposes TriRoute, a single lightweight per-token, per-layer controller that jointly chooses (i) attention resolution (skip/local/full), (ii) a sparse FFN expert set including a null expert that recovers MoD-style skips, and (iii) KV-cache bit-width, under one Lagrangian multi-resource budget. Training uses a heterogeneous relaxation (ST-Gumbel for attention/bits, load-balanced top-k for experts), per-axis gradient balancing, whitening, and a hinge entropy floor to arrest a described cross-axis collapse cascade. The central claim is that, on 160M–1.3B decoder-only models at compute-optimal token counts, this joint controller Pareto-dominates the best independently tuned MoD+MoE+KV-quantization combination at matched inference FLOPs and memory and better preserves tail robustness (rare entities, code, math, GSM8K), with an interpretable policy that spends read/write fidelity on rare and boundary tokens while decoupling FFN compute.

Significance. If the joint-routing advantage is real, the contribution is substantial: it reframes three largely separate conditional-computation literatures as one coupled allocation problem and supplies a concrete controller, cost model, and balancing recipe that practitioners could adopt. The design principles (separate read vs write decisions; null expert as MoD special case) and the collapse-cascade diagnosis are useful even without new SOTA numbers. Credit is due for a fully specified method (§3), an explicit multi-resource cost model (Eqs. 8–9, Appendix A), a reference PyTorch controller (Appendix C), and a clear honesty note that the evaluation protocol is intended to be reproducible. The work is therefore significant as a methods unification; its significance as an empirical Pareto result depends entirely on whether the reported tables are measured end-to-end outcomes.

major comments (3)
  1. Load-bearing empirical support is missing as written. The Abstract, §5.1 (Table 3, Fig. 4), and §5.3 (Table 5) state that TriRoute Pareto-dominates the independent MoD+MoE+KV-quant combo and preserves tail metrics on completed 160M–1.3B compute-optimal runs (including seed variance <0.05 and A100 throughput). §4’s honesty note and Appendix D state the opposite: numbers “illustrate the trends our design targets,” “should be regenerated from full training runs … before use as benchmark claims,” and several figures are “illustrative.” Without measured runs, the dominance and tail-robustness claims have no empirical support. Either report full end-to-end results under the stated protocol, or reframe the paper as a methods contribution and remove/qualify all performance claims in the Abstract and §5.
  2. §4 / Table 3 baseline protocol is under-specified for a claim of “best independently-tuned combination.” The independent combo is said to be grid-searched to the joint (0.55, 0.40) budget, but the search space, number of trials, and whether MoD’s block gate, MoE sparsity, and KV bit-width were co-tuned or only separately swept then combined are not given. Once real runs exist, this must be documented so the Pareto gap cannot be attributed to a weak multi-axis baseline rather than joint learning.
  3. §5.2 / Fig. 3: the cross-axis collapse cascade is central to the training recipe, yet the figure is labeled “illustrative of the qualitative dynamics,” and no quantitative entropy/load trajectories from actual joint training are reported. The necessity of whitening (Eq. 11), gradient rescaling (Eq. 7), and the entropy floor (Eq. 12) is asserted via ablations in Table 4 that inherit the same non-measured status. Provide real training curves and ablations, or demote these components from “essential” empirical findings to proposed stabilizers pending validation.

Circularity Check

0 steps flagged

No significant circularity: joint-routing claims are architectural/empirical, not forced by definition, fit, or self-citation.

full rationale

TriRoute’s load-bearing chain is (i) a shared controller emitting attention mode, expert gates (with null expert), and KV bit-width (Eqs. 2–5), (ii) heterogeneous ST-Gumbel / top-k relaxations with per-axis gradient balancing (Eqs. 6–7), (iii) a differentiable multi-resource cost plus Lagrangian budget (Eqs. 8–14), and (iv) coupling-aware whitening/entropy-floor losses to arrest a collapse cascade (Eqs. 11–12). None of these steps defines the quality metric in terms of the claimed outcome: the LM loss is independent of the budget duals, so Pareto dominance over independently tuned MoD+MoE+KV-quant at matched FLOPs/memory is not true by construction and must be measured. There is no fitted scalar that is then re-reported as a prediction; dual ascent only enforces the cost constraint. References are external (Switch/GShard, MoD, KIVI/KVQuant, Gumbel-Softmax, etc.); the authors do not import a uniqueness theorem or ansatz from their own prior work. Recovering MoD via a null expert is an acknowledged special case, not a renaming of a known result as a new derivation. Side-features (position, surprisal, previous decisions) create ordinary closed-loop training, not a definitional loop. Appendix D’s caveat that tables are protocol illustrations is an empirical-honesty issue, not circularity of the derivation. Score 0 with empty steps is therefore the correct finding.

Axiom & Free-Parameter Ledger

7 free parameters · 5 axioms · 3 invented entities

The paper is an empirical systems/methods contribution. It inherits the standard transformer stack and discrete-routing estimators from prior work, then adds many hand-chosen training knobs and a few named constructs (shared TriRoute controller, null expert as MoD special case, cross-axis collapse cascade, Lagrangian multi-resource budget). Free parameters dominate: balancing weights, temperatures, expert counts, bit set, budget targets, and dual step sizes are chosen rather than derived. No new physical entity is postulated; invented entities are architectural.

free parameters (7)
  • Budget targets C*_flops, C*_mem = (0.55, 0.40) default
    Primary operating point fixed at (0.55, 0.40) of dense cost; sweeps define the frontier. Chosen by authors, not derived.
  • Balancing weights α, β, γ and entropy floor ζ = α=1e-2, β=1e-3, γ=1e-3, ζ=0.5
    Load-balance, z-loss, entropy-floor coefficients and ζ=0.5 control collapse prevention; single setting transferred across scales by hand.
  • Gumbel temperatures τ_a, τ_b anneal schedule = 2.0 → 0.5 cosine
    Annealed 2.0→0.5; controls bias-variance of categorical relaxations.
  • Expert count E, top-k, null expert, router width d_r = E=8, top-k=2, d_r=128
    Architecture capacity and controller size chosen so active FFN FLOPs match dense; d_r=128 claimed <1% FLOPs.
  • Attention mode set and local window w; bit set B = w=128; bits {2,4,8,16}
    Discrete option sets {skip,local,full} and {2,4,8,16} with w=128 are design choices that bound the policy space.
  • Dual ascent step ρ_λ and EMA momentum m = ρ_λ=0.05, m=0.99
    Online Lagrangian and gradient/whitening EMAs are hand-tuned controllers for budget and scale matching.
  • Peak LR and token budgets by scale = LR 3–6e-4; 3.2B/8.2B/26B tokens
    Training schedule scaled with model size under compute-optimal ~20 tokens/parameter assumption from Hoffmann et al.
axioms (5)
  • domain assumption Straight-through Gumbel-Softmax and top-k softmax provide usable biased gradients for joint discrete routing of heterogeneous axes.
    Invoked throughout §3.4; standard in discrete routing literature but not proved to yield near-optimal hard policies after annealing.
  • ad hoc to paper A single shared trunk representation of token importance can usefully inform attention, expert, and bit decisions without fatal interference at the studied sparsities.
    Design Principle and §5.2 sharing ablations; partially supported only by the paper’s own (illustrative) ablations.
  • domain assumption Differentiable expected FLOPs/memory costs (Eq. 8–9, App. A) match executed hard-decision cost closely enough that Lagrangian training hits the intended operating point.
    §3.5 and §3.8; authors note residual router overhead and imperfect kernels in §7.
  • domain assumption Compute-optimal token counts (~20 tokens/parameter) and Pythia-matched decoder-only backbones are a fair testbed for comparing allocation methods.
    §4 Table 2; standard scaling-law practice but limits claims about larger or heavily overtrained regimes.
  • ad hoc to paper Per-axis whitening, gradient-norm rescaling, and hinge entropy floors arrest cross-axis collapse under a shared trunk.
    §3.6 and Figure 3; presented as the fix for the cascade the authors identify.
invented entities (3)
  • TriRoute shared multi-axis controller no independent evidence
    purpose: Emit coordinated per-token policies over attention mode, experts, and KV bits under one budget.
    Core architectural proposal; independent evidence would be third-party replications or production deployments, not provided here.
  • Cross-axis routing-collapse cascade no independent evidence
    purpose: Name and motivate the failure mode where collapse on one routing axis induces collapse on the others via a shared trunk.
    Empirical phenomenon claimed in §3.6/Figure 3; figure labeled illustrative, so the entity is not yet independently measured.
  • Null expert as MoD-style FFN skip inside joint routing no independent evidence
    purpose: Recover depth-skip as a special case while allowing attention and bits to stay high.
    Special case of MoE/MoD ideas; the joint use with separate attention/bit heads is the paper-specific construct.

pith-pipeline@v1.1.0-grok45 · 23858 in / 4537 out tokens · 49892 ms · 2026-07-11T03:19:53.585091+00:00 · methodology

0 comments
read the original abstract

Conditional computation can decouple language model quality from per-token inference cost, yet leading techniques act on a single axis in isolation: Mixture-of-Experts (MoE) sparsifies the FFN, Mixture-of-Depths (MoD) skips whole transformer blocks, and KV-cache quantization compresses attention memory. We argue these three decisions (attention resolution, expert selection, and cache bit-width) are strongly coupled and should be made jointly: a token rare enough to warrant full attention may also need high-precision caching regardless of which expert processes it. We introduce TriRoute, a single lightweight controller shared across all three axes that, for every token at every layer, emits a coordinated policy: (i) an attention mode (skip/local/full), (ii) a sparse set of FFN experts (with a null expert recovering MoD), and (iii) a KV-cache bit-width. The controller trains end-to-end via a heterogeneous relaxation (Gumbel-Softmax with straight-through estimation for categorical decisions and load-balanced top-k gating for experts) under a Lagrangian budget constraint that turns the average compute and memory cost into a controllable knob. We identify a cross-axis routing-collapse cascade in naive joint training, where collapse on one axis propagates to the others, and address it with per-axis normalization and a coupling-aware balancing loss. On decoder-only models from 160M to 1.3B parameters at compute-optimal token counts, TriRoute Pareto-dominates the best independent MoD+MoE+KV-quantization combination at matched inference FLOPs and memory, while better preserving tail-case robustness on rare entities, code, and arithmetic that pure perplexity optimization erodes. Post-hoc analysis reveals interpretable structure: the controller allocates full attention and high-precision cache to sentence-initial positions, rare subwords, and named entities, while cheaply routing function words.

Figures

Figures reproduced from arXiv: 2607.06601 by Andrii Balashov, Olena Ponomarova.

Figure 1
Figure 1. Figure 1: From three isolated mechanisms to one controller. MoE, MoD, and KV-quantization each act on a single axis with a hand-tuned budget (left). TriRoute (right) replaces them with a shared per-token controller gϕ that emits a coupled policy over attention resolution, expert selection, and KV precision under one global budget. Tokens are drawn with area proportional to the compute+memory they receive: the rare e… view at source ↗
Figure 2
Figure 2. Figure 2: A TriRoute block. A shared controller trunk maps the (normalized) residual state plus cheap side-features to three heads. The attention head picks a query mode (skip/local/full) governing how much of the past the token reads; the expert head selects top-k of E FFN experts or the null expert (FFN skip); the bit head sets the precision at which the token’s own K/V are written to the cache for future tokens t… view at source ↗
Figure 3
Figure 3. Figure 3: The cross-axis collapse cascade and its fix. Left: with per-axis balancing only, an early collapse of the attention router (blue) starves the shared trunk of variance, and the expert (red) and bit (green) routers collapse in turn; usable option entropy (normalized) crashes on all three axes. Right: adding per-axis feature whitening (Equation (11)) and the marginal-entropy floor keeps all three routers dive… view at source ↗
Figure 4
Figure 4. Figure 4: TriRoute Pareto-dominates the independent combination on both the compute (a) and memory (b) frontiers (1.3B). Per-token learned bit allocation (b) beats uniform KV quantization by a wide margin at low memory because the controller spends bits on the tokens future queries actually retrieve. The dense model (⋆) is matched at ∼55% FLOPs / 40% cache. 6 What Does the Controller Learn? A unified controller is o… view at source ↗
Figure 5
Figure 5. Figure 5: Learned policy by token category and depth (1.3B). Each cell is the mean resource a category receives at a depth band. The controller keeps attention and cache bits high for sentence￾initial tokens, rare/entity tokens, and numerals, while FFN compute concentrates on content words, numbers, and code. Rare entities show the signature high-attention, high-bits, low-FFN pattern that a block-level method cannot… view at source ↗

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

56 extracted references · 56 canonical work pages · 17 internal anchors

  1. [1]

    GQA: Training generalized multi-query transformer models from multi-head checkpoints

    Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. GQA: Training generalized multi-query transformer models from multi-head checkpoints. InEmpirical Methods in Natural Language Processing (EMNLP), 2023

    work page 2023

  2. [2]

    CoLT5: Faster long-range transformers with conditional computation.Empirical Methods in Natural Language Processing (EMNLP), 2023

    Joshua Ainslie, Tao Lei, Michiel de Jong, Santiago Ontañón, Siddhartha Brahma, Yury Zemlyanskiy, David Uthus, Mandy Guo, James Lee-Thorp, Yi Tay, Yun-Hsuan Sung, and Sumit Sanghai. CoLT5: Faster long-range transformers with conditional computation.Empirical Methods in Natural Language Processing (EMNLP), 2023

    work page 2023

  3. [3]

    Longformer: The Long-Document Transformer

    Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020

    work page internal anchor Pith review Pith/arXiv arXiv 2004

  4. [4]

    Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation

    Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation.arXiv preprint arXiv:1308.3432, 2013. 16 TriRoute: Unified Adaptive-Compute Routing

    work page internal anchor Pith review Pith/arXiv arXiv 2013

  5. [5]

    Pythia: A suite for analyzing large language models across training and scaling

    Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O’Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning (ICML), 2023

    work page 2023

  6. [6]

    PIQA: Reasoning about physical commonsense in natural language

    Yonatan Bisk, Rowan Zellers, Jianfeng Gao, and Yejin Choi. PIQA: Reasoning about physical commonsense in natural language. InAAAI Conference on Artificial Intelligence (AAAI), 2020

    work page 2020

  7. [7]

    Generating Long Sequences with Sparse Transformers

    Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers.arXiv preprint arXiv:1904.10509, 2019

    work page internal anchor Pith review Pith/arXiv arXiv 1904

  8. [8]

    Unified scaling laws for routed language models

    Aidan Clark, Diego de Las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann, Bogdan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, et al. Unified scaling laws for routed language models. InInternational Conference on Machine Learning (ICML), 2022

    work page 2022

  9. [9]

    Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge

    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

    work page internal anchor Pith review Pith/arXiv arXiv 2018

  10. [10]

    Training Verifiers to Solve Math Word Problems

    Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems.arXiv preprint arXiv:2110.14168, 2021

    work page internal anchor Pith review Pith/arXiv arXiv 2021

  11. [11]

    SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention

    Róbert Csordás, Piotr Piękos, Kazuki Irie, and Jürgen Schmidhuber. SwitchHead: Accelerating transformers with mixture-of-experts attention.arXiv preprint arXiv:2312.07987, 2024

    work page internal anchor Pith review Pith/arXiv arXiv 2024

  12. [12]

    Damai Dai, Chengqi Deng, Chenggang Zhao, R. X. Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y. Wu, et al. DeepSeekMoE: Towards ultimate expert specialization in mixture-of-experts language models.arXiv preprint arXiv:2401.06066, 2024

    work page internal anchor Pith review Pith/arXiv arXiv 2024

  13. [13]

    Fu, Stefano Ermon, Atri Rudra, and Christopher Ré

    Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. InAdvances in Neural Information Processing Systems (NeurIPS), 2022

    work page 2022

  14. [14]

    DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model

    DeepSeek-AI. DeepSeek-V2: A strong, economical, and efficient mixture-of-experts language model.arXiv preprint arXiv:2405.04434, 2024

    work page internal anchor Pith review Pith/arXiv arXiv 2024

  15. [15]

    Uni- versal transformers

    Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Uni- versal transformers. InInternational Conference on Learning Representations (ICLR), 2019

    work page 2019

  16. [16]

    LLM.int8(): 8-bit matrix multiplication for transformers at scale

    Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. LLM.int8(): 8-bit matrix multiplication for transformers at scale. InAdvances in Neural Information Processing Systems (NeurIPS), 2022

    work page 2022

  17. [17]

    QLoRA: Efficient finetuningofquantizedLLMs

    Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. QLoRA: Efficient finetuningofquantizedLLMs. InAdvances in Neural Information Processing Systems (NeurIPS), 2023

    work page 2023

  18. [18]

    Depth-adaptive transformer

    Maha Elbayad, Jiatao Gu, Edouard Grave, and Michael Auli. Depth-adaptive transformer. In International Conference on Learning Representations (ICLR), 2020. 17 TriRoute: Unified Adaptive-Compute Routing

    work page 2020

  19. [19]

    Reducing transformer depth on demand with structured dropout

    Angela Fan, Edouard Grave, and Armand Joulin. Reducing transformer depth on demand with structured dropout. InInternational Conference on Learning Representations (ICLR), 2020

    work page 2020

  20. [20]

    Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity.Journal of Machine Learning Research (JMLR), 23(120):1–39, 2022

    William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity.Journal of Machine Learning Research (JMLR), 23(120):1–39, 2022

    work page 2022

  21. [21]

    GPTQ: Accurate post- training quantization for generative pre-trained transformers

    Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. GPTQ: Accurate post- training quantization for generative pre-trained transformers. InInternational Conference on Learning Representations (ICLR), 2023

    work page 2023

  22. [22]

    MegaBlocks: Efficient sparse training with mixture-of-experts

    Trevor Gale, Deepak Narayanan, Cliff Young, and Matei Zaharia. MegaBlocks: Efficient sparse training with mixture-of-experts. InMachine Learning and Systems (MLSys), 2023

    work page 2023

  23. [23]

    The Pile: An 800GB Dataset of Diverse Text for Language Modeling

    Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800GB dataset of diverse text for language modeling.arXiv preprint arXiv:2101.00027, 2020

    work page internal anchor Pith review Pith/arXiv arXiv 2020

  24. [24]

    Adaptive Computation Time for Recurrent Neural Networks

    Alex Graves. Adaptive computation time for recurrent neural networks.arXiv preprint arXiv:1603.08983, 2016

    work page internal anchor Pith review Pith/arXiv arXiv 2016

  25. [25]

    Training Compute-Optimal Large Language Models

    Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models.arXiv preprint arXiv:2203.15556, 2022

    work page internal anchor Pith review Pith/arXiv arXiv 2022

  26. [26]

    Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami

    Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Michael W. Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami. KVQuant: Towards 10 million context length LLM inference with KV cache quantization. InAdvances in Neural Information Processing Systems (NeurIPS), 2024

    work page 2024

  27. [27]

    Categorical reparameterization with Gumbel-Softmax

    Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with Gumbel-Softmax. InInternational Conference on Learning Representations (ICLR), 2017

    work page 2017

  28. [28]

    Mixtral of Experts

    Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts.arXiv preprint arXiv:2401.04088, 2024

    work page internal anchor Pith review Pith/arXiv arXiv 2024

  29. [29]

    Scaling Laws for Neural Language Models

    Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models.arXiv preprint arXiv:2001.08361, 2020

    work page internal anchor Pith review Pith/arXiv arXiv 2001

  30. [30]

    Reformer: The efficient transformer

    Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations (ICLR), 2020

    work page 2020

  31. [31]

    GShard: Scaling giant models with conditional computation and automatic sharding

    Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. GShard: Scaling giant models with conditional computation and automatic sharding. InInternational Conference on Learning Representations (ICLR), 2021

    work page 2021

  32. [32]

    BASE layers: Simplifying training of large, sparse models

    Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. BASE layers: Simplifying training of large, sparse models. InInternational Conference on Machine Learning (ICML), 2021. 18 TriRoute: Unified Adaptive-Compute Routing

    work page 2021

  33. [33]

    KIVI: A tuning-free asymmetric 2bit quantization for KV cache

    Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. KIVI: A tuning-free asymmetric 2bit quantization for KV cache. In International Conference on Machine Learning (ICML), 2024

    work page 2024

  34. [34]

    Decoupled weight decay regularization

    Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. InInternational Conference on Learning Representations (ICLR), 2019

    work page 2019

  35. [35]

    Maddison, Andriy Mnih, and Yee Whye Teh

    Chris J. Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxationofdiscreterandomvariables. InInternational Conference on Learning Representations (ICLR), 2017

    work page 2017

  36. [36]

    The LAMBADA dataset: Word prediction requiring a broad discourse context

    Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The LAMBADA dataset: Word prediction requiring a broad discourse context. InAssociation for Computational Linguistics (ACL), 2016

    work page 2016

  37. [37]

    DeepSpeed-MoE: Advancing mixture-of- experts inference and training to power next-generation AI scale

    Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Ammar Ahmad Awan, Jeff Rasley, and Yuxiong He. DeepSpeed-MoE: Advancing mixture-of- experts inference and training to power next-generation AI scale. InInternational Conference on Machine Learning (ICML), 2022

    work page 2022

  38. [38]

    Mixture-of-Depths: Dynamically allocating compute in transformer-based language models

    David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, and Adam Santoro. Mixture-of-depths: Dynamically allocating compute in transformer-based language models.arXiv preprint arXiv:2404.02258, 2024

    work page internal anchor Pith review Pith/arXiv arXiv 2024

  39. [39]

    Hash layers for large sparse models

    Stephen Roller, Sainbayar Sukhbaatar, Arthur Szlam, and Jason Weston. Hash layers for large sparse models. InAdvances in Neural Information Processing Systems (NeurIPS), 2021

    work page 2021

  40. [40]

    WinoGrande: An adversarial winograd schema challenge at scale.Communications of the ACM, 64(9):99–106, 2021

    Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. WinoGrande: An adversarial winograd schema challenge at scale.Communications of the ACM, 64(9):99–106, 2021

    work page 2021

  41. [41]

    Confident adaptive language modeling

    Tal Schuster, Adam Fisch, Jai Gupta, Mostafa Dehghani, Dara Bahri, Vinh Tran, Yi Tay, and Donald Metzler. Confident adaptive language modeling. InAdvances in Neural Information Processing Systems (NeurIPS), 2022

    work page 2022

  42. [42]

    Fast Transformer Decoding: One Write-Head is All You Need

    Noam Shazeer. Fast transformer decoding: One write-head is all you need.arXiv preprint arXiv:1911.02150, 2019

    work page internal anchor Pith review Pith/arXiv arXiv 1911

  43. [43]

    GLU Variants Improve Transformer

    Noam Shazeer. GLU variants improve transformer.arXiv preprint arXiv:2002.05202, 2020

    work page internal anchor Pith review Pith/arXiv arXiv 2002

  44. [44]

    Outrageously large neural networks: The sparsely-gated mixture-of-experts layer

    Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. InInternational Conference on Learning Representations (ICLR), 2017

    work page 2017

  45. [45]

    RoFormer: Enhanced transformer with rotary position embedding.Neurocomputing, 568:127063, 2024

    Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. RoFormer: Enhanced transformer with rotary position embedding.Neurocomputing, 568:127063, 2024

    work page 2024

  46. [46]

    Adaptive attention span in transformers

    Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive attention span in transformers. InAssociation for Computational Linguistics (ACL), 2019

    work page 2019

  47. [47]

    RedPajama: An open source recipe to reproduce LLaMA training dataset

    Together Computer. RedPajama: An open source recipe to reproduce LLaMA training dataset. https://github.com/togethercomputer/RedPajama-Data, 2023. 19 TriRoute: Unified Adaptive-Compute Routing

    work page 2023

  48. [48]

    Williams

    Ronald J. Williams. Simple statistical gradient-following algorithms for connectionist reinforce- ment learning.Machine Learning, 8(3):229–256, 1992

    work page 1992

  49. [49]

    Efficient streaming language models with attention sinks

    Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. InInternational Conference on Learning Representations (ICLR), 2024

    work page 2024

  50. [50]

    Big bird: Transformers for longer sequences

    Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences. InAdvances in Neural Information Processing Systems (NeurIPS), 2020

    work page 2020

  51. [51]

    HellaSwag: Can a machine really finish your sentence? InAssociation for Computational Linguistics (ACL), 2019

    Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. HellaSwag: Can a machine really finish your sentence? InAssociation for Computational Linguistics (ACL), 2019

    work page 2019

  52. [52]

    Root mean square layer normalization

    Biao Zhang and Rico Sennrich. Root mean square layer normalization. InAdvances in Neural Information Processing Systems (NeurIPS), 2019

    work page 2019

  53. [53]

    Mixture of attention heads: Selecting attention heads per token

    Xiaofeng Zhang, Yikang Shen, Zeyu Huang, Jie Zhou, Wenge Rong, and Zhang Xiong. Mixture of attention heads: Selecting attention heads per token. InEmpirical Methods in Natural Language Processing (EMNLP), 2022

    work page 2022

  54. [54]

    H2O: Heavy-hitter oracle for efficient generative inference of large language models

    Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, Zhangyang Wang, and Beidi Chen. H2O: Heavy-hitter oracle for efficient generative inference of large language models. InAdvances in Neural Information Processing Systems (NeurIPS), 2023

    work page 2023

  55. [55]

    Dai, Zhifeng Chen, Quoc V

    Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M. Dai, Zhifeng Chen, Quoc V. Le, and James Laudon. Mixture-of-experts with expert choice routing. InAdvances in Neural Information Processing Systems (NeurIPS), 2022

    work page 2022

  56. [56]

    ST-MoE: Designing Stable and Transferable Sparse Expert Models

    Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. ST-MoE: Designing stable and transferable sparse expert models.arXiv preprint arXiv:2202.08906, 2022. 20 TriRoute: Unified Adaptive-Compute Routing A Extended Cost Model We give the normalization used to turn raw per-token costs into the budget fract...

    work page internal anchor Pith review Pith/arXiv arXiv 2022