cs.LGNov 7, 2025

Controllably Efficient Language Models

Authors: Jatin PrakashAahlad PuliRajesh Ranganath

Organizations: New York University

Abstract

The substantial inference costs of attention in transformers motivated the development of efficient sequence mixers: namely sparse and sliding window attention, convolutions and linear attention. Although these approaches result in impressive reductions in inference costs, they often trade-off with quality, specifically in-context recall. Apriori fixing this quality-cost tradeoff at training time means being suboptimal from the get-go: some downstream applications might fundamentally require more memory for in-context recall, while other tasks may require lower latency and memory. We propose a conceptually simple meta-sequence mixer with inference-cost controllability: the Compress & Attend Transformer (CAT). CAT decodes chunks of tokens by attending to compressed chunks of the sequence so far. Both compression and decoding can use any existing sequence mixer. Decoding from the compressed sequence yields compute and memory savings, with chunk size setting the operating point on the quality-cost trade-off. Importantly, training CAT across multiple chunk sizes at once unlocks test-time control of this trade-off without any retraining, all in a single model. Instantiated with the most basic choice, dense attention as the mixer, CAT surprisingly suffices to match 10 popular and diverse efficient models (linear, hybrids, sparse) on real-world long-context recall at comparable inference costs, all from a single trained model. CAT further performs competitively on long-context understanding benchmarks while providing 1.4-3.7x higher generation throughput than a dense transformer. Code is at: https://github.com/rajesh-lab/cat-transformer

Explore similar work

Apr 24, 2026cs.LG

Toeplitz MLP Mixers are Low Complexity, Information-Rich Sequence Models

Transformer-based large language models are in some respects limited by the quadratic time and space computational complexity of attention. We introduce the Toeplitz MLP Mixer (TMM), a transformer-like architecture that swaps attention for triangular-masked Toeplitz matrix multiplication over the sequence dimension resulting in O(dnlogn)\mathcal{O} (dn \log n) time and O(dn)\mathcal O(dn) space complexity during training and O(dn)\mathcal O(dn) time and space at inference prefill. Despite the lack of sophisticated input modulation or state maintenance present in other sub-quadratic architectures, TMMs yield greater training efficiency in terms of loss achieved per compute and device memory. We demonstrate that TMMs are capable of retaining more input information resulting in improved copying ability, which we argue results from a lack of architectural biases. Consistent with higher input information retention, TMMs exhibit superior information retrieval and in-context learning benchmark accuracy compared to comparable architectures. We conclude with an analysis from the perspective of operator index theory and show that, counterintuitively, trained Toeplitz layers of causal non-invertible models are more likely to be invertible or nearly so than models that are actually invertible over their inputs.
Benjamin L. Badger, Ethan Roland
Jun 22, 2026cs.LG

SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers

Long contexts have become standard in pretrained LLMs, yet they remain expensive to run: prefill compute grows quadratically with sequence length, and every decode step re-reads a key-value cache that grows linearly with it. Sparse attention cuts these costs by attending only to a relevant subset of past tokens, but selecting that subset is itself expensive. We present SpotAttention, a lightweight selector that attaches to a frozen pretrained transformer and learns by KL distillation to estimate its attention distribution. The selector picks the top-K keys each query attends to, and because its estimate is a calibrated distribution, a dual top-p rule reads the per-query, per-layer budget directly from it. Across Qwen3 (dense, 4B-32B) and Qwen3.5 (hybrid linear/full attention, 4B-9B), SpotAttention matches dense accuracy at contexts up to 128K tokens, eight times the training length. Decode at L=128K runs 3.9x faster than FlashAttention and 1.8x faster than Twilight, the strongest training-free baseline. Quantizing the selector's K-cache to INT4 or FP4 microscale shrinks it 3.5x at no accuracy cost.
Huzama Ahmad, Se-Young Yun
May 27, 2026cs.LG

Meta-Attention: Bayesian Per-Token Routing for Efficient Transformer Inference

Standard transformer architectures apply a single attention mechanism uniformly across all tokens and sequence positions, irrespective of local context or computational budget. We propose Meta-Attention, a framework that dynamically routes each token to the most appropriate attention strategy -- full softmax attention, linear (kernel) attention, or sliding-window local attention -- via a Bayesian Meta-Controller. Unlike prior routing approaches that use deterministic or prior-free learned routing, the Meta-Controller treats per-token mechanism selection as posterior inference under a compute-aware Dirichlet prior: routing weights are the output of an amortised variational posterior q(alpha | x_t; phi) trained with an Evidence Lower Bound (ELBO) objective that jointly encodes task performance and attention-mechanism cost. This design produces principled routing uncertainty estimates that govern the soft-to-hard routing transition, mitigates routing collapse without ad hoc load-balancing losses, and yields better compute-performance trade-offs than deterministic or prior-free learned routing at negligible overhead. Phase 1 empirical results on a Tiny LM benchmark confirm core predictions: the Bayesian controller's learned routing distribution implies a projected normalised FLOP cost of 25.1% under hard routing, vs. 59.3% for the prior-free baseline (-34.2 pp), and reduces routing entropy from 55.8% to 43.3% (-12.5 pp), demonstrating that the Dirichlet prior prevents routing collapse while the non-Bayesian model defaults to full attention. We present the Bayesian architecture, ELBO training objective, and a Phase 1 PyTorch prototype validating forward-pass correctness, posterior diversity, and a controlled ablation against a prior-free baseline. Code available at: https://github.com/KFEAL/meta-attention
Alan Ferrari