PARTREP: Learning What to Repeat for Decoder-only LLMs

• URL: http://arxiv.org/abs/2607.01792v1
• Date: Thu, 02 Jul 2026 07:07:28 GMT
• Title: PARTREP: Learning What to Repeat for Decoder-only LLMs
• Abstract summary: PartRep is a selective augmentation method that appends only the most informative tokens.<n>We train a lightweight gate that predicts high-NLL tokens from early-layer hidden states.<n>PartRep retains most of the gains of full repetition while using only 59.4% of its KV cache and 79.0% of its prefill FLOPs.
• Score: 12.053018926036678
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: While decoder-only LLMs excel at a vast array of natural language tasks, it suffers from an asymmetric information flow induced by causal attention: later tokens are richer in contextual grounding than earlier ones. A simple and effective remedy is prompt repetition -- just appending a second copy of prompt before generation can redistribute grounding across positions and improve reasoning performance. However, full repetition of the original prompt doubles the KV cache footprint and quadruples attention cost during prefill, making it impractical for long-context settings. We propose PartRep, a selective augmentation method that appends only the most informative tokens -- rather than the entire prompt. We use token-wise negative log-likelihood (NLL) as a selection signal, motivated by the hypothesis that less predictable tokens are less recoverable from surrounding context and therefore benefit more from late-position repetition. To avoid the heavy cost of a full forward pass for scoring, we train a lightweight gate that predicts high-NLL tokens from early-layer hidden states, enabling token selection during mid-prefill via early exit. Across eight benchmarks (including MMLU, GSM8K, and RULER) and three model families (Qwen2.5, Llama3.2, Gemma4), PartRep retains most of the gains of full repetition while using only 59.4\% of its KV cache and 79.0\% of its prefill FLOPs.

Related papers

• Roll Out and Roll Back: Diffusion LLMs are Their Own Efficiency Teachers [76.15132587294862]
  Wide-In, Narrow-Out (WINO) is a training-free decoding algorithm that enables revokable parallel generation.<n>WINO+ injects the verified denoising trajectories produced by WINO into model parameters, aligning training with efficient inference.<n>Experiments on LLaDA and MMaDA show that WINO improves both quality and efficiency, while WINO+ further strengthens this progression.
  arXiv  Detail & Related papers  (2026-05-16T11:27:40Z)
• RCP: Representation Consistency Pruner for Mitigating Distribution Shift in Large Vision-Language Models [24.3914653184824]
  Large Vision-Language Models (LVLMs) suffer from prohibitive inference costs due to the massive number of visual tokens processed by the language decoder.<n>Existing pruning methods often lead to significant performance degradation because the irreversible removal of visual tokens causes a distribution shift in the hidden states that deviates from the pre-trained full-token regime.<n>We propose Consistency Representation Pruner, which integrates cumulative visual token pruning with a delayed repair mechanism.
  arXiv  Detail & Related papers  (2026-04-04T13:31:45Z)
• Behavior-Equivalent Token: Single-Token Replacement for Long Prompts in LLMs [55.827877498548965]
  We propose a lightweight training framework that learns a single prompt-specific Behavior-Equivalent token ([BE])<n>The framework first trains [BE] to encode the natural-language content of the original system prompt via reconstruction, and then distills the prompt's downstream behavior into this single token.<n> Empirical evaluations on three datasets show that a single [BE] token achieves up to a 3000x reduction in prompt length, while retaining about 98% of the downstream performance of the original system prompts.
  arXiv  Detail & Related papers  (2025-11-28T15:22:52Z)
• Think-at-Hard: Selective Latent Iterations to Improve Reasoning Language Models [22.525318796588568]
  Think-at-Hard (TaH) is a dynamic latent thinking method that iterates deeper only at hard tokens.<n>TaH boosts reasoning performance across five challenging benchmarks.
  arXiv  Detail & Related papers  (2025-11-11T18:57:02Z)
• FrugalPrompt: Reducing Contextual Overhead in Large Language Models via Token Attribution [3.4666771782038652]
  Large language models (LLMs) owe much of their stellar performance to expansive input contexts, yet such verbosity inflates monetary costs, carbon footprint, and inference-time latency.<n>We introduce FrugalPrompt, a novel prompt compression framework for LLMs, which retains only the most semantically significant tokens.<n>We evaluate the approach across four NLP tasks: Sentiment Analysis, Commonsense QA, Summarization, and Mathematical Reasoning.
  arXiv  Detail & Related papers  (2025-10-18T10:22:13Z)
• Multipole Attention for Efficient Long Context Reasoning [64.94673641704289]
  Large Reasoning Models (LRMs) have shown promising accuracy improvements on complex problem-solving tasks.<n>LRMs need to generate long chain-of-thought reasoning in order to think before answering.<n>We introduce Multipole Attention, which accelerates autoregressive reasoning by only computing exact attention for the most important tokens.
  arXiv  Detail & Related papers  (2025-06-16T03:00:40Z)
• Token Assorted: Mixing Latent and Text Tokens for Improved Language Model Reasoning [53.57895922042783]
  Large Language Models (LLMs) excel at reasoning and planning when trained on chainof-thought (CoT) data.<n>We propose a hybrid representation of the reasoning process, where we partially abstract away the initial reasoning steps using latent discrete tokens.
  arXiv  Detail & Related papers  (2025-02-05T15:33:00Z)
• An Early FIRST Reproduction and Improvements to Single-Token Decoding for Fast Listwise Reranking [50.81324768683995]
  FIRST is a novel approach that integrates a learning-to-rank objective and leveraging the logits of only the first generated token. We extend the evaluation of FIRST to the TREC Deep Learning datasets (DL19-22), validating its robustness across diverse domains. Our experiments confirm that fast reranking with single-token logits does not compromise out-of-domain reranking quality.
  arXiv  Detail & Related papers  (2024-11-08T12:08:17Z)
• Mitigating the Learning Bias towards Repetition by Self-Contrastive Training for Open-Ended Generation [92.42032403795879]
  We show that pretrained language models (LMs) such as GPT2 still tend to generate repetitive texts. We attribute their overestimation of token-level repetition probabilities to the learning bias. We find that LMs use longer-range dependencies to predict repetitive tokens than non-repetitive ones, which may be the cause of sentence-level repetition loops.
  arXiv  Detail & Related papers  (2023-07-04T07:53:55Z)
• RetroMAE v2: Duplex Masked Auto-Encoder For Pre-Training Retrieval-Oriented Language Models [3.4523793651427113]
  We propose duplex masked auto-encoder, a.k.a. DupMAE, which targets on improving the semantic representation capacity for contextualized embeddings of both [] and ordinary tokens. DupMAE is simple but empirically competitive: with a small decoding cost, it substantially contributes to the model's representation capability and transferability.
  arXiv  Detail & Related papers  (2022-11-16T08:57:55Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.