PARTREP: Learning What to Repeat for Decoder-only LLMs

2026-07-02

AI Summary

This work addresses the inefficient utilization of context in decoder-only large language models caused by causal attention, where full-prompt repetition improves performance but incurs prohibitive computational and memory costs. The authors propose PartRep, a novel approach that leverages token-level negative log-likelihood (NLL) as a metric of information importance. By predicting high-NLL tokens using hidden states from early layers and integrating lightweight gating with an early-exit strategy, PartRep dynamically selects critical tokens for repetition during the prefill phase. Evaluated across eight benchmarks and three model families, this method achieves performance nearly on par with full repetition while substantially reducing KV cache usage to 59.4% and prefill computation to 21.0% of the original cost.

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.

Problem

• asymmetric information flow
• prompt repetition
• KV cache
• prefill cost
• decoder-only LLMs

Innovation

• selective repetition
• negative log-likelihood
• early exit
• KV cache optimization
• decoder-only LLMs