Research map / Architectures and efficiency / Attention and transformers

Distilling Sequential Computation in Transformer Language Models

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Efficient language transformers · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes UMIM, a lightweight attention-pooling merge module that replaces frequent n-gram token spans with single surrogate embeddings (distilled from a frozen LM), reducing effective sequence length by up to ~40% at inference with rollback-based KV-cache merging.

Problem

Autoregressive Transformer LMs process every token separately, making long prompts and growing decoding contexts expensive; the paper asks whether multi-token computation can be replaced at inference by single-step input surrogates without updating the backbone.

Approach

A shared single-layer multi-head attention-pooling merge module maps frequency-mined n-gram spans (n=2-4) to one static surrogate embedding, trained by KL distillation between the frozen LM's predictions on original vs. merged inputs; at decoding time matched suffixes roll back KV entries and substitute one merged entry.

Key results

• Llama-3.1-8B merge module reaches 37.8% token reduction with Top-1 alignment 79.2% and MRR 87.1% against the original LLM.
• On WikiText-103 with Llama-3-8B, the merge module keeps PPL at 13.8 (original 13.4) at 36.1% token reduction, while SelectContext and LLMLingua-2 reach 65.0 and 76.5 PPL.
• Task-Adapted merge module on ARC-Easy reaches 77.88% accuracy at 22% token reduction, exceeding the 76.49% uncompressed baseline.

Surprising findings

• Task-adapted merging on ARC-Easy and HellaSwag exceeds the uncompressed backbone accuracy at 20-40% token reduction.
• Early attempts with cosine-similarity-based averaging of adjacent token embeddings failed to preserve model behavior.

Limitations

• Only bigram-level task rules are re-mined in the task-adaptation experiments, capping token reduction at 50%.
• GPT2-XL shows larger PPL degradation after merging; context-dependent or task-critical spans may incur greater approximation error.
• Study limited to a few backbone sizes and domains; future work on adaptive merging and broader scales.

Experimental setup

Datasets: WikiText-103 WikiText-103 derived merge rules long reasoning trajectory corpus (self-collected) HellaSwag training split ARC training split

Benchmarks: WikiText-103 BookCorpus OpenWebText PIQA Copa OpenBookQA ARC-Easy ARC-Challenge CNN/DailyMail AIME 2024 AMC HellaSwag

Metrics: perplexity token reduction accuracy top-1 overlap top-3 overlap top-p overlap MRR ROUGE BERTScore pass@k

Baselines: SelectContext LLMLingua LLMLingua-2 H2O StreamingLLM CPC EHPC DeepScaleR

Builds on: Zip2Zip LLMLingua-2 Select Context

Closest papers by meaning

• Bigram Subnetworks: Mapping to Next Tokens in Transformer Language Models (NeurIPS 2025 · similarity 0.90)
• Transformers need glasses! Information over-squashing in language tasks (NeurIPS 2024 · similarity 0.89)
• Compress, Gather, and Recompute: REFORMing Long-Context Processing in Transformers (NeurIPS 2025 · similarity 0.89)
• Effective Context in Transformers: An Analysis of Fragmentation and Tokenization (NeurIPS 2026 · similarity 0.89)
• Pause Tokens Strictly Increase the Expressivity of Constant-Depth Transformers (NeurIPS 2025 · similarity 0.89)
• Separations in the Representational Capabilities of Transformers and Recurrent Architectures (NeurIPS 2024 · similarity 0.88)