Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

PARTREP: Learning What to Repeat for Decoder-only LLMs

About

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.

Andikawati P Widjaja, Yongjun Kim, Hyounghun Kim, Jaeho Lee• 2026

Related benchmarks

TaskDatasetResultRank
General KnowledgeMMLU (test)
Accuracy63.2
71
Science Question AnsweringARC Challenge (test)
Accuracy81.4
60
Question AnsweringSciQ (test)
Accuracy92.2
36
Long-context evaluationRULER 4k--
35
KnowledgeOpenBookQA (test)
Accuracy77.8
19
Complex ReasoningMMLU Pro (test)
Accuracy27.8
8
Complex ReasoningGSM8K (test)
Accuracy85.5
8
General knowledge and scientific retrievalMedQA (test)
Accuracy50.3
8
Long-context evaluationRULER 8K context length
Accuracy83.8
4
Long-context evaluationRULER 16k context length
Accuracy73.2
4
Showing 10 of 10 rows

Other info

Follow for update