PaperlayerView exact revision on arXiv

Source-linked AI summary

Distilling Sequential Computation in Transformer Language Models

Zixuan Lan, Jessica Yang, Yanhong Li, Karen Livescu, Jiawei Zhou

arXiv:2609.27233v1cs.CL

TL;DR

Long Transformer contexts are costly, and it is unclear whether multi-step computation can be replaced at inference by single-step surrogates without updating the pretrained backbone. UMIM merges token spans into surrogate embeddings and applies them to prompts and decoding contexts, reducing effective sequence length by up to 40% with minimal accuracy degradation.

  • Problem

    Autoregressive Transformers process tokens sequentially, making long inputs and growing contexts expensive, while it remains unclear whether multi-step computation can be replaced at inference without internal activations or backbone updates.

    Source: §1 · pp. 1–2

  • Method

    UMIM uses a lightweight shared module to compute a single surrogate embedding from a selected span’s static token embeddings, replacing the span in pretrained-model inputs.

    Source: §1 · p. 2

  • Results

    Up to 40% effective sequence-length reduction with minimal accuracy degradation is reported across diverse models and tasks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    UMIM demonstrates input-level compression for frozen language models across representative models and tasks without changing the backbone.

    Source: §5 · p. 10

    The study focuses on representative models and tasks; broader model scales and domains remain future work.

    Source: Limitations · p. 11

Abstract

from arXiv · show

Transformer language models process sequences token by token in an autoregressive manner, making growing contexts increasingly expensive. Yet many adjacent token spans are highly predictable or frequently occur as stable units, suggesting that their representations may be compressible. We introduce a method for distilling sequential computation by replacing spans of input tokens with collapsed representations, computed on the fly by a lightweight merge module. This module generates a single surrogate embedding from a sequence of static token embeddings that captures the functional role of the multiple tokens, allowing pretrained models to operate on compressed inputs without architectural changes or re-training. We apply this approach during inference to compress both prompts and intermediate decoding steps, using a rollback mechanism to substitute stored multi-token KV cache entries with their single-step surrogates. Experiments across diverse models show that the merge module can be used to reduce effective sequence length by up to 40% with minimal accuracy degradation across language modeling evaluations and downstream tasks, including question answering, summarization, commonsense reasoning, and long-form mathematical reasoning. Additional lightweight adaptation of the merge module further improves the accuracy-compression trade-off in selected settings. These results demonstrate that sequential token computation in Transformers can be effectively approximated through condensed surrogate representations that approximate the original behavior without model updating.

1 Introduction

Long Transformer contexts are costly, and it remains unclear whether multi-step computation can be replaced at inference with single-step surrogates without backbone updates or internal activations. The paper introduces UMIM to merge token spans into surrogate embeddings and reports up to 40% effective sequence-length reduction with minimal accuracy degradation.

1 Introduction

  • Autoregressive models process tokens sequentially, making long inputs and growing contexts expensive despite frequent, predictable adjacent spans.“Machine learning” and “New York” are examples of spans that often occur together.Source: §1 · p. 1
  • The open question is whether multi-step Transformer computation can be replaced at inference by single-step surrogates without internal activations or pretrained-backbone updates.Existing approaches include alternative tokenization, multi-token decoding, and context reduction, often with retraining, auxiliary components, or trade-offs.Source: §1 · pp. 1–2
  • UMIM computes one surrogate embedding on demand from the static embeddings of a selected token span, with the same n-gram receiving the same surrogate across contexts.Different n-grams receive different surrogates.Source: §1 · p. 2
  • The shared module is a lightweight plug-in trained with predictive distillation while the pretrained language model remains frozen.Task-specific merge rules and lightweight module adaptation can further improve downstream performance.Source: §1 · p. 2
  • Up to 40% effective sequence-length reduction with minimal accuracy degradation is reported across diverse models and tasks using UMIM.The method compresses both prompts and autoregressive decoding contexts while preserving next-token predictive behavior.Source: §1 · p. 2

2 Related Work

Related work reduces sequential computation through changes to tokenization, learned context representations, prompt compression, and KV-cache or token management. UMIM is distinguished as post-training span merging at the input level, without changing the tokenizer or deleting, filtering, or rewriting content.

Tokenization Variations and Extension

  • Tokenization and chunk-decoding approaches use alternative input units or extended span representations, but generally require tokenizer or token-space changes, auxiliary components, or retraining.Examples include byte-level models such as ByT5 [Xue et al., 2022], larger phrase-level units such as SuperBPE [Liu et al., 2025], and chunk-level methods such as CoG [Lan et al., 2023].Source: §2 · p. 2

Latent Context Compression

  • Learned latent-context methods encode inputs into compact vectors or special-token representations, typically through an additional encoding stage and a separate compression interface.Examples include Gist Tokens [Mu et al., 2023] and ICAE [Ge et al., 2024].Source: §2 · p. 2

Prompt Compression and Context Reduction

  • Prompt-compression methods retain task-relevant content, filter less informative segments, or rewrite prompts; UMIM is task-agnostic and does not delete, filter, or rewrite input content.The passage describes task-aware approaches such as LongLLMLingua [Jiang et al., 2024] and task-agnostic approaches such as LLMLingua [Jiang et al., 2023].Source: §2 · pp. 2–3

KV Cache Management, Token Pruning and Merging

  • KV-cache management and token-pruning or merging methods act on cache entries or hidden states, whereas UMIM merges spans before Transformer computation and substitutes single-step surrogate states.The comparison covers cache management and token merging such as MRT5 [Kallini et al., 2024].Source: §2 · p. 3

3 Methodology

UMIM maps selected contiguous token spans to single surrogate embeddings computed from static input embeddings, then applies those surrogates during prompt processing and autoregressive decoding. Its method combines corpus-derived merge rules, a lightweight attention-pooling module, predictive distillation, and runtime KV-cache rollback.

3 Methodology

  • UMIM is an external plug-in that maps selected contiguous token spans to surrogate embeddings, replacing the original embeddings while aiming to preserve subsequent predictive distributions.The merge module operates on input embeddings before Transformer computation.Source: §3 · p. 3

3.1 Merge Rule

  • Merge rules select corpus-frequency n-grams of lengths 2–4 above threshold τ, resolve overlaps with longest match, and can use in-domain data to improve coverage.The rule set depends on tokenizer and corpus statistics, not model architecture.Source: §3.1 · p. 3

3.2 Lightweight Merge Module Architecture

  • The lightweight module uses learnable-query, multi-head attention pooling over static span embeddings to produce one surrogate in the model’s embedding space.A shared key-value projection and an output projection are used in the pooling architecture.Source: §3.2 · pp. 3–4

3.3 UMIM Learning by Distilling Sequential computation

  • Distillation minimizes KL divergence between predictions on unmerged and merged prefixes while keeping the language model frozen; a merged span collapses its token steps to one retained position.Figure 1 depicts the token-step-aligned distillation loss and merge-module architecture.Source: §3.3 · p. 4, Figure 1 · p. 4

3.4 Model Inference with Runtime Merging

  • At inference, UMIM replaces eligible prompt spans and checks each generated suffix, rolling back matched KV entries into one merged state; reduced sequence length lowers cache size and attention computation.The checking and rollback overhead is amortized over long sequences; merge rules may also be adapted with in-domain data.Source: §3.4 · pp. 4–5, Figure 2 · p. 5

4 Experiments

Experiments assess predictive alignment, language modeling, generation, and downstream transfer, alongside task-specific adaptation, ablations, and inference efficiency. Results show compression across settings, with adaptation improving accuracy–compression trade-offs on evaluated tasks.

4.1 Experimental Setup

  • The experimental setup spans language-model backbones, distribution-alignment metrics, perplexity and token reduction, and downstream QA, summarization, and reasoning evaluations.The task-agnostic setting applies a WikiText-103-trained module without task-specific tuning; the adaptation setting specializes the merge rules and module while keeping the backbone frozen.Source: §4.1 · p. 5

4.2 Predictive Distribution Alignment

  • 35.7%–54.0% effective sequence-length reduction accompanies strong predictive-distribution alignment across backbones.Table 1 reports alignment metrics in percentages; higher values indicate better agreement with the original model’s predictive distributions at aligned positions.Source: Table 1 · p. 6, §4.2 · p. 6

4.3 Language Modeling and Generation

  • Nearly 40% of tokens are removed on WikiText-103, while BookCorpus and OpenWebText show 12–15% token reduction without retraining.Llama-3.1-8B remains close to its original model across the three datasets; GPT2-XL has a larger perplexity increase but remains stronger than external prompt-compression baselines at comparable compression rates.Source: §4.3 · p. 6, Table 2 · p. 6

4.4 Task-Agnostic Generalization to Downstream Tasks

  • Across unseen QA tasks and three backbones, UMIM provides the strongest or most robust accuracy–compression trade-off; on CNN/DailyMail it achieves the best Rouge and BERTScore among compressed variants.The same WikiText-103-trained module transfers without task-specific tuning; in long-form math reasoning, it remains more robust than Select Context and LLMLingua2 above 40% token reduction.Source: §4.4 · p. 7

4.5 Task-Specific Adaptation of a Pretrained Merge Module

  • 84.86% accuracy at a 21.94% merge rate, and 81.06% at 39.53%, for Task-Adapted on HellaSwag; it achieves the strongest accuracy–compression trade-off.At 39.53% merge rate, accuracy remains 2.72 percentage points above the uncompressed baseline; Figure 3 also shows it outperforming external baselines at comparable merge rates.Source: §4.5 · p. 9, Figure 3 · p. 9
  • 77.88% accuracy at 22.00% token reduction on ARC-Easy, and 54.85% accuracy at 20.00% token reduction on ARC-Challenge, after adaptation.These results exceed the uncompressed ARC-Easy baseline of 76.49% and ARC-Challenge baseline of 49.50%.Source: §4.5 · p. 9

4.6 Ablation Study and Efficiency Analysis

  • Ablations find rule design relatively robust when merges provide sufficient training signal, but overly aggressive merging degrades performance; larger module capacity brings only marginal gains.The current attention-based design performs well with a small parameter budget, while simpler linear variants and deeper, less efficient structures do not improve the trade-off.Source: §4.6 · pp. 9–10
  • UMIM improves end-to-end prefill and decoding throughput at a fixed 50% merge ratio, with larger prefill gains for longer input contexts.The study reports tokens per second on Llama-3.1-8B across settings denoted by prefill length plus decode length.Source: §4.6 · p. 10

5 Conclusion

UMIM replaces matched multi-token spans with surrogate embeddings while leaving the pretrained backbone unchanged. It compresses prompts and decoding contexts, transfers to unseen corpora and downstream tasks without retraining, and can be adapted for improved accuracy–compression trade-offs.

5 Conclusion

  • UMIM maps matched multi-token spans to one surrogate embedding rather than deleting tokens or changing the Transformer architecture.The merge module operates on static input embeddings while the backbone language model remains unchanged.Source: §5 · p. 10
  • At inference, UMIM merges eligible prompt spans before prefill and replaces matching suffix KV entries after generating each token, shortening later sequence and cache use without skipping generation.Source: §5 · p. 10
  • The WikiText-trained merge module transfers directly to unseen corpora and downstream tasks without retraining.Source: §5 · p. 10
  • For a new task, re-mining merge rules increases coverage, while fine-tuning followed by DPO further improves the accuracy–compression trade-off with the backbone frozen.Source: §5 · p. 10

Limitations

The study establishes UMIM as an input-level compression framework for frozen language models across representative models and tasks. Broader scales and domains, adaptive merging, and interactions with downstream adaptation remain for future work.

Limitations

  • The study focuses on establishing input-level compression for frozen language models across representative models and tasks.Source: Limitations · p. 11
  • Broader model scales and domains are identified as areas for future work.Source: Limitations · p. 11
  • Future work may explore adaptive merging strategies and interactions between pretrained merge modules and downstream adaptation.Source: Limitations · p. 11

Technical Appendices

The appendices trace UMIM from failed similarity heuristics to a learned input-level surrogate, explain its approximation scope, and give a conditional account of generalization across contexts.

A.1 Empirical Summary

  • Task-agnostic transfer and lightweight task adaptation preserve strong accuracy–compression trade-offs, with adaptation supporting higher merge ratios and sometimes improving downstream accuracy.The WikiText-103-trained module transfers without task-specific supervision; adaptation updates only the merge module while freezing the backbone.Source: §A.1 · p. 18

A.2 LLM Background

  • Because autoregressive generation processes tokens sequentially, long inputs and generated contexts increase inference cost, motivating compression that preserves predictive behavior.Source: §A.2 · p. 18

A.3 Early Attempt One: Embedding Similarity-Based Token Merging

  • Embedding-similarity merging failed because semantic proximity did not ensure computational substitutability, while small threshold changes could cause excessive merging or negligible compression.The appendix then motivates learning a surrogate from the frozen model’s predictive behavior instead of averaging similar embeddings.Source: §A.3 · p. 19, §A.4 · p. 19

A.4 Early Attempt 2: Training-Based Pairwise Token Merging

  • A frozen GPT-2 experiment showed pairwise merged representations could approximate next-token distributions, but scaling to IMDB exposed stability, capacity, and generalization challenges.The fixed segmentation and over-parameterized network were empirically chosen and motivated a more scalable formulation.Source: §A.4 · pp. 19–20

A.5 Foundation to our current Work

  • Frequency-based candidate spans and predictive distillation train a shared surrogate to approximate a recurring span’s functional effect across observed contexts, rather than reconstructing its token embeddings or hidden states.The appendices motivate frequent spans as better supported by repeated context observations and frame UMIM as complementing, not replacing, BPE segmentation.Source: §A.5 · pp. 20–21

A.6 Information Preservation and Approximation Scope

  • UMIM targets predictive behavior rather than exact internal-state recovery, but ambiguous, context-dependent, or task-critical spans can incur greater approximation error.Future queries cannot access every constituent token independently after multiple KV states are compressed into one surrogate.Source: §A.6 · pp. 21–22

A.7 Generalization of Context-Shared Surrogate Embeddings

  • The generalization analysis separates irreducible surrogate-approximation error from finite-context statistical uncertainty, which decreases as minimum frequency support increases.Its guarantee is conditional on fitting observed teacher behavior and applies to unseen contexts drawn from the same distribution, not arbitrary spans or contexts.Source: §A.7 · pp. 22–23

A.8 UMIM as Parameter-Efficient Task Adaptation

UMIM reuses or adapts a lightweight merge module for downstream tasks while keeping the backbone frozen. Its current 2–4-gram operating range is a practical choice, and the conclusion presents learned span surrogates as an approximation rather than an exact replacement.

A.8 UMIM as Parameter-Efficient Task Adaptation

  • Across three evaluated tasks, adapting only the merge module supports substantially higher merge ratios while preserving and sometimes improving downstream accuracy.Task adaptation reuses the WikiText-trained initialization; the backbone remains frozen, while task-specific rules, distillation-based SFT, and optional RL adapt the module.Source: §A.8 · pp. 23–24

A.9 Choice of n-gram Lengths

  • The 2–4-gram range is a practical balance among frequency support, rule-set size, training feasibility, and earlier compression during decoding, not an architectural upper bound.Longer spans may have fewer observations and become eligible for merging later; the module can process variable-length spans.Source: §A.9 · pp. 24–25

A.10 Conclusion

  • The conclusion frames learned input-level surrogates as empirical evidence that recurring-span computation can be compressed while approximately preserving prediction, without claiming exact equivalence in every context.The authors also identify limited training data and computational resources as current constraints, and propose context-dependent or hierarchical surrogates and long-context training as future directions.Source: §A.10 · p. 25

B Merge-Rule Mining Details

The paper builds merge rules from frequent token n-grams, then filters overlaps and assigns non-overlapping matches. Rule sources and adaptation vary by setting, while 2- to 4-gram rules are a practical design choice.

Rule construction

  • The base rules retain contiguous 2- to 4-grams occurring at least five times, with the threshold varied to trade merge coverage against statistical support.The tokenized corpus is segmented before n-gram counting; counts are aggregated across segments.Source: B Merge-Rule Mining Details · pp. 25–26

Rule filtering

  • Same-order overlapping rules keep the more frequent candidate, and lower-order rules contained in retained higher-order spans are removed.This filtering prioritizes longer, more specific spans when the same local pattern occurs at multiple orders.Source: B Merge-Rule Mining Details · p. 26

Match assignment

  • Matches are ordered by starting position and then span length, and retained only when they do not overlap; merging is single-stage, not recursive.The resulting locations are determined before the forward pass.Source: B Merge-Rule Mining Details · p. 26

Task-specific rules

  • Task adaptation mines rules from the target task’s training split; WT + Rules reconstructs rules, while WT + FT and Task-Adapted also adapt the WikiText-trained module.Base task-agnostic experiments instead use WikiText-103 rules and the WikiText-trained module without downstream-task data for rule reconstruction.Source: B Merge-Rule Mining Details · p. 26

Rule scope

  • Using 2- to 4-grams is a practical choice reflecting recurring-span support, longer-rule coverage, and decoding delay, not an architectural limit.For base training, Table 7 reports retained rule counts at frequency threshold τ = 5; Llama-3.1-8B and Llama-3.2-1B share a rule set, while DeepScaleR uses a generated mathematical-reasoning corpus.Source: B Merge-Rule Mining Details · p. 26, Table 7 · p. 27

C Module Architecture and Integration

UMIM pools selected token spans into backbone-compatible surrogate embeddings while leaving the Transformer frozen. The module supports training-time, prompt-time, and decoding-time compression without modifying internal Transformer layers.

C Module Architecture and Integration

  • The attention-based pooling module maps selected static token embeddings to one surrogate embedding with the backbone’s input dimension; current experiments use spans of two to four tokens.The pooling architecture can process variable-length spans.Source: C Module Architecture and Integration · p. 27

C.1 Attention-based Pooling Architecture

  • Each head uses a learned, input-independent query, while attention weights remain input-dependent through the projected span, producing span-specific pooled representations.Head outputs are concatenated and passed through an output processor to form the surrogate embedding.Source: §C.1 · p. 27

C.2 Difference from Standard Self-Attention

  • Unlike standard self-attention, the module uses one learned query per head to reduce a whole span to one vector rather than returning an output sequence of the same length.The shared key-value projection keeps the module compact; separate projections or extra layers did not clearly improve results, and the selected design balanced alignment, stability, and parameter efficiency.Source: §C.2 · p. 28

C.3 Parameter Count

  • For Llama-3.1-8B, the 33.6-million-parameter merge module is about 0.42% of the backbone, and only module parameters are updated.The module is trained separately per backbone; direct checkpoint transfer across backbones is not assumed.Source: C Module Architecture and Integration · p. 27, §C.3 · p. 28

C.4 Training-time Integration

  • During training, selected spans are replaced by surrogates, and compressed predictions are aligned with frozen-backbone teacher predictions using soft-target cross-entropy.The objective is equivalent to forward KL divergence up to the teacher distribution’s fixed entropy; the backbone remains unchanged, and equal-length spans can be pooled in parallel.Source: §C.4 · p. 28

C.5 Prompt-time Integration

  • At prompt time, non-overlapping rule matches are replaced at the embedding level before an ordinary backbone pass, yielding a shorter prompt KV cache.Matching proceeds left to right, preferring longer spans at the same start; positions in the compressed sequence are contiguous.Source: §C.5 · p. 28

C.6 Decoding-time Integration

  • During decoding, UMIM checks each generated-token suffix against rules and, on a match, rolls back cached suffix states and processes the span as one surrogate input.Generated token IDs and decoded text are unchanged; the operation shortens the effective sequence and KV-cache representation, without recursively merging prior surrogates.Source: §C.6 · p. 29

D Predictive Distillation Loss and Alignment Metrics

This section describes how UMIM trains its merge module by aligning teacher and compressed-path predictions, distinguishes distillation from alignment metrics, and applies learned surrogates to prompts and decoding-time KV caches. It also covers inference-time cache and computation trade-offs and task adaptation of merge rules and the merge module while keeping the backbone frozen.

D Predictive Distillation Loss and Alignment Metrics

  • Aligned teacher and student predictions exclude the unaffected prefix and pair each merged span at its final original position with its surrogate position.Length-preserving placeholders place the surrogate at the span’s final position before inactive entries are removed and the compressed sequence is packed in causal order.Source: D Predictive Distillation Loss and… · p. 30
  • Full-vocabulary soft-target cross-entropy trains only the merge module and has the same gradient as forward KL because the teacher and backbone are frozen.The loss is normalized over aligned positions rather than padded sequence length.Source: D Predictive Distillation Loss and… · p. 31
  • Full-vocabulary targets preserve the teacher’s full distribution, while tested Top-k targets lacked reliable convergence; full-vocabulary training costs more memory and computation, but only during training.The Top-k runs were exploratory, not a controlled quantitative ablation; Top-k and Top-p are alignment metrics, not training objectives.Source: D Predictive Distillation Loss and… · p. 31
  • Predictive agreement is evaluated, not trained, with Top-k overlap at k = 3 and 10, Top-p overlap, and mean reciprocal rank on aligned positions.Padding is added after position selection, and padded rows are excluded from both loss and metrics.Source: D Predictive Distillation Loss and… · p. 32

E Inference with UMIM: Prompt and Decoding-Time Merging

  • Prompt merging scans left to right, selecting longest matching spans first, and replaces each with a surrogate before prefill, reducing prompt and initial KV-cache length.With decoding-time merging disabled, generation then proceeds normally; the experiments check 4-, 3-, and 2-gram candidates.Source: E Inference with UMIM: Prompt and… · p. 33
  • When a generated token completes a matching suffix, UMIM crops the preceding raw-token KV entries and forwards a surrogate to create a new KV state, without averaging the original entries.Merges are longest-first and do not overlap or nest; one token is still selected at every autoregressive step, reducing later effective sequence and cache lengths rather than skipping generation steps.Source: E Inference with UMIM: Prompt and… · pp. 33–35
  • Decoding-time merging reduces internal sequence and KV-cache lengths without removing generated tokens, but cannot skip autoregressive steps and adds lookup, merge, and cache-cropping work.These costs can make realized decoding speedup smaller than prefill speedup even when cache length is reduced.Source: E Inference with UMIM: Prompt and… · p. 36

F Task Adaptation with SFT and RL

  • Task adaptation differs from task-agnostic evaluation by rebuilding merge rules from target-task training data and optionally updating the merge module through SFT and RL.The task-agnostic setting instead uses WikiText-103 rules and its trained base merge module without downstream training data.Source: F Task Adaptation with SFT and RL · p. 36

F.1 Task-Specific Merge-Rule Construction

  • Task-specific rules retain adjacent bigrams meeting a frequency threshold and apply them using left-to-right, non-overlapping matching.Lower thresholds retain more spans and can increase merge ratio, while higher thresholds produce more conservative rule sets.Source: §F.1 · p. 37

F.2 Distillation-Based SFT

  • Distillation-based SFT trains only the merge module by aligning frozen-backbone predictions on original and compressed task sequences with soft-target cross-entropy.For multiple-choice tasks such as HellaSwag, SFT sequences pair the context with the correct ending.Source: §F.2 · p. 37

F.3 RL Adaptation

  • Task adaptation combines predictive-distillation SFT with RL to optimize task-level decisions through the merge module while keeping the backbone frozen.SFT preserves general predictive behavior under the task distribution; RL then encourages higher scores for correct answers.Source: §F.3 · p. 38

F.4 Adaptation and Ablation Settings

  • Task-trained uses a randomly initialized merge module, while WT + DPO skips SFT; both are ablations rather than the final task-adaptation pipeline.Task-trained measures the benefit of WikiText-trained initialization, and WT + DPO applies RL directly to the WikiText-trained base module.Source: §F.4 · p. 38

F.5 Task-Adaptation Experimental Configuration

Task adaptation changes task-specific merge rules and updates the merge module while keeping the backbone frozen. The configuration uses distillation-based SFT followed by a separate DPO stage, with task-specific data and selection procedures.

F.5 Task-Adaptation Experimental Configuration

  • All adaptation experiments use Llama-3.1-8B; the tokenizer, input embeddings, and backbone stay frozen while only merge-module parameters are updated.Except for Task-trained, the module starts from the WikiText-103-trained checkpoint.Source: §F.5 · p. 38
  • HellaSwag rules are training-split bigrams with frequency thresholds from 2000 to 1, yielding realized merge rates from 12.10% to 43.09%.ARC-Easy and ARC-Challenge use threshold-5 bigrams mined from training questions; validation examples guide model selection but do not construct rules or update parameters.Source: §F.5 · p. 39
  • SFT distills teacher next-token distributions from the frozen uncompressed model into predictions from the compressed model, updating only the merge module.WT + FT starts from the WikiText checkpoint, whereas Task-trained starts from a newly initialized module; SFT checkpoints use validation distillation loss and early stopping.Source: §F.5 · p. 39
  • The DPO stage uses correct continuations as preferred and incorrect continuations as rejected, optimizing the policy module’s relative preference while the backbone and reference module remain frozen.The policy and reference start from the same SFT checkpoint; WT + DPO instead initializes directly from the WikiText module, and RL checkpoints use validation accuracy.Source: §F.5 · p. 39

F.6 Implementation Procedures

  • SFT and RL run as separate stages connected by the merge-module checkpoint, with RL further optimizing task-level decisions after SFT adapts the module.The backbone remains frozen throughout, and the two objectives are not jointly optimized.Source: §F.6 · pp. 40, 42

G Base Merge-Module Training Setup

  • The base merge-module training setup is described separately from task adaptation, which is covered in Section F.This section identifies corpora, optimization settings, hardware, and model-specific configurations as its setup topics.Source: G Base Merge-Module Training Setup · p. 42

G.1 Training Corpora

The section describes training UMIM across general-language and math-reasoning corpora, then evaluates its likelihood accounting, generation behavior, and downstream use. Results vary by task and decoding strategy, with strong reported performance for prompt compression in summarization and long-form mathematical reasoning.

G.1 Training Corpora

  • WikiText-103 supplies training data and merge rules for the Llama and GPT-2 modules, while a separate generated reasoning corpus trains DeepScaleR’s math module.The math corpus uses sampled reasoning trajectories, concatenated and divided into 1,024-token sequences; splits are 90%/5%/5%.Source: §G.1 · p. 42

G.2 Optimization and Distributed Training

  • Training compares compressed and original predictions through a frozen backbone; only merge-module parameters are updated, with distillation cross-entropy optimized and other agreement metrics monitored.Each module uses four pooling heads, with dimensions determined by its backbone embedding space.Source: §G.2 · p. 43, §G.3 · p. 43

H Additional Experiments and Evaluation

  • The additional evaluations include task-agnostic likelihood transfer, generation, and multiple-choice QA using WikiText-trained components, with no task-specific QA rules or tuning.BookCorpus and OpenWebText test transfer beyond the merge module’s WikiText training corpus.Source: §H.1 · p. 43, §H.3.1 · pp. 46–47

H.1 Perplexity Evaluation

  • For perplexity, a two-stage procedure scores transitions between compressed elements and tokens skipped inside merged spans, then normalizes by original next-token targets.This token-level likelihood accounting is distinct from decoding-time KV-cache rollback.Source: §H.1 · pp. 43–45

H.2.2 Qualitative Analysis

  • Qualitative examples show fluent continuations and continued attention to surrogate states, but neither attention nor near-zero cosine similarity establishes semantic equivalence.Merged embeddings are distinct from copying or averaging constituent embeddings.Source: §H.2.2 · pp. 45–46

H.2.3 Quantitative Results

  • Generation outcomes depend on decoding: Top-p improves diversity and reduces repetition, while unrestricted random sampling yields mixed similarity, diversity, and repetition results.Greedy decoding slightly reduces repetition, keeps most metrics close, and yields a higher Mauve score.Source: §H.2.3 · p. 46

H.3.2 CNN/DailyMail Summarization

  • The fixed WikiText-trained UMIM module remains competitive with and outperforms evaluated baselines on CNN/DailyMail at comparable token-reduction levels.The experiment applies merging to the prompt only, without summarization-specific training.Source: §H.3.2 · p. 47

H.3.3 Long-Form Mathematical Reasoning

  • UMIM achieves strong prompt-compression performance on long-form mathematical reasoning, outperforming Selective Context and LLMLingua in the reported results.The evaluation uses DeepScaleR-1.5B-Preview, allows up to 15,000 generated tokens, and reports pass@16.Source: §H.3.3 · pp. 47–48

H.3.4 Long-Context Evaluation

The long-context evaluations test retrieval, multi-document QA, and summarization using a frozen WikiText-trained merge module, while a separate experiment tests code-adapted UMIM. Performance varies with context length and task, and code results retain a gap to the uncompressed baseline.

H.3.4 Long-Context Evaluation

  • Long-context evaluation uses a frozen WikiText-trained module and backbone, re-mining eligible n-gram rules for evaluation corpora because the original rules provide too few merges.The tasks include four-needle retrieval, multi-document QA, and long-document summarization; merging is used in prompt and decoding for summarization.Source: §H.3.4 · pp. 48–49
  • UMIM stays close to the uncompressed baseline on shorter retrieval and multi-document QA contexts, but performance gaps grow with longer contexts; summarization quality is comparable on GovReport and lower on QMSum.GovReport also shows slight improvements on several metrics while reducing prompt length and generated-token KV-cache usage.Source: §H.3.4 · p. 49, Table 16 · p. 49, Table 17 · p. 49
  • The long-context results transfer a merge module trained on 512-token WikiText-103 sequences to retrieval and QA contexts up to 16K tokens and documents up to 51K tokens, without long-context training.The results also show that training sequence length affects preservation of long-range dependencies.Source: §H.3.4 · p. 49

H.4 Specialized-Domain Evaluation: Code-Adapted UMIM

  • At 25% token reduction, code-adapted UMIM retains approximately 81% of original HumanEval pass@1 performance, from 62.2% to 50.6%.LLMLingua-2 and Selective Context retain substantially less performance at the same reduction level.Source: §H.4 · p. 50
  • Code generation remains challenging: the HumanEval gap to the uncompressed baseline may improve with larger in-domain corpora or task adaptation.The code experiment uses domain-specific merge rules and training data, and is presented as supplementary validation given limited data and compute.Source: §H.4 · p. 50

I.1 Empirical Efficiency Evaluation and Throughput Analysis

The efficiency evaluation measures prefill and decoding throughput separately under a controlled 50% merge ratio. UMIM substantially improves prefill throughput, while decoding gains are smaller; the current Python-level prototype incurs overhead.

I.1 Empirical Efficiency Evaluation and Throughput Analysis

  • The throughput tests use Llama-3.1-8B on one A100 GPU at batch size 1, with generation fixed at 256 tokens and prompt length varied from 256 to 2,048.Each setting uses 10 warm-up runs followed by 10 timed runs.Source: §I.1 · pp. 50–51
  • Prefill speedup rises from approximately 1.46× to 3.74× as prompt length increases from 256 to 2,048 tokens, while decoding throughput improves by 0.9%–2.5%.Throughput is measured against original token counts under a controlled 50% merge ratio.Source: §I.1 · pp. 50–51, Table 19 · p. 51
  • Prefill benefits because all matching prompt spans can be merged before the backbone runs, shortening the sequence throughout the complete prefill pass.During decoding, merging shortens the KV cache used by later steps but does not bypass autoregressive generation steps.Source: §I.1 · p. 51
  • The non-fused Python-level prototype adds scheduling and kernel-launch overhead, offsetting part of the shorter-cache benefit; results indicate initial feasibility, not a fully optimized system.The reported positive decoding throughput gains are achieved despite this overhead.Source: §I.1 · pp. 51–52

J Ablation Study

The ablations examine merge-rule selection, module capacity, and positional handling. They find that frequency trades candidate coverage against statistical support, extra depth adds little in the tested setting, and compressed-position merging works empirically without establishing exact equivalence.

J.1 Frequency Threshold and Statistical Support

  • Top-1 rises from 0.7410 at τ = 5 to 0.8437 at τ = 500 on Llama-3.2-1B, while MRR rises from 0.8297 to 0.9078; higher thresholds improve alignment but reduce candidate coverage.Lower thresholds admit more spans but include candidates with less training support.Source: §J.1 · p. 52

J.2 Merge-Module Capacity

  • Top-1 increases from 0.7410 to 0.7427 when depth grows from one to three attention layers, and the remaining metrics change by at most approximately 0.0011.The authors use the lighter one-layer design in the main experiments.Source: §J.2 · pp. 52–53, Table 21 · p. 53

J.3 Alternative Merge-Rule Construction

  • Confidence-based rules use 2- to 4-gram spans whose token probabilities exceed 0.5, and achieve strong distribution-preservation metrics at a realized TR of 10.3%.Because token reduction is modest, the result is a feasibility study, not a matched-TR comparison with frequency-based rules.Source: §J.3 · p. 53, Table 22 · p. 53

J Ablation Study

  • The ablations show that frequency trades candidate coverage against statistical support, extra module depth adds little in the evaluated setting, and alternative rule criteria can also work.High-frequency n-grams are an effective implementation of UMIM, not a fixed framework limitation.Source: §J.3 · p. 53

J.4 Positional Encoding under Token-Span Merging

  • Training and inference use the same compressed-position convention, allowing distillation to account jointly for surrogate replacement and position compaction without claiming exact equivalence.The surrogate is a learned continuous soft token intended to preserve predictive behavior, not embedding or intermediate-state equivalence.Source: §J.4 · p. 54
  • With RoPE, relative offsets remain unchanged within regions that do not cross merged spans, while distances across merged spans shorten; this local account is not an end-to-end error bound.UMIM also shows empirical viability with learned absolute positional embeddings, though position shifts can still matter.Source: §J.4 · p. 55
  • Experiments do not isolate error from surrogate replacement versus position compaction, while adapting only the merge module can further reduce predictive discrepancy with the backbone frozen.The complete method is reported to work across multiple models, tasks, and positional-encoding architectures.Source: §J.4 · pp. 55–56
  • Preserving original position indices is an unevaluated alternative; dynamic decoding would require tracking logical and physical positions and changing cache and position bookkeeping.The paper leaves systematic implementation and evaluation of this alternative positional design to future work.Source: §J.4 · p. 56

Related papers

Loading 2609.27233v1…