EntroPrefill: Renyi-Guided Context Pruning with Conditional Stability Guarantees for Retrieval-Augmented Generation

What it claims

EntroPrefill is a mid-prefill pruning procedure for RAG. At a selected layer it drops whole retrieved chunks and runs the remaining layers only on the retained tokens. It keeps the shallow KV entries and original position IDs, so the model does not re-run on a shortened prompt. The contribution is not attention-guided pruning, which the paper attributes to LLMLingua, AttentionRAG, LazyLLM, SlimInfer and ASL. It is the coupling of an entropy-based proposal with three separate checks:

• hard per-row caps on discarded attention mass
• a statistically valid audit
• a conditional output bound

The paper separates three kinds of guarantee:

• exact, on the proposal rows
• probabilistic, over a fixed query-observer distribution
• conditional, for the full-model output

There is no measured speedup or accuracy result.

How it works

1. Protection. Instructions, the full user query, delimiters and initial sink tokens are never removed. A chunk that overlaps protected tokens is protected whole. All documents must precede the query.
2. Head pooling (Rényi order 2). For each head, attention averaged over the proposal query rows is renormalized excluding sink positions. Head weight is proportional to exp(−H₂/τ), which is exponential in negative entropy and not 1/entropy. It is mixed with a floor: w = (1−ζ)v + ζ/r within each GQA group. The floor stops a diffuse or sink-dominated head from being zeroed out. With a positive floor, the sink head's weight tends to ζ/r. Heads with too little non-sink mass are ineligible. If every group falls back, pruning is bypassed.
3. Chunk score. S_i = ln ℓ_i + λκ_i: - ℓ_i is the chunk's smoothed attention mass relative to its token share. - κ_i ∈ [0,1] is within-chunk concentration. - Concentration alone can never overturn a relative-mass ratio larger than e^λ.
4. Constrained greedy removal. Chunks are sorted by ascending score. Each is added only if every per-row discarded-mass cap ϵ_a and a minimum retained-chunk count m₀ still hold. Cost is O(k log k) plus O(k·H_Q·W_p). This is a heuristic. - A dual bound U_l(y,t) (Eq. 15) upper-bounds the best achievable token saving for any nonnegative prices. It can prove a token budget infeasible and bounds the greedy's suboptimality. - It cannot prove a budget feasible.
5. Audit and layer selection. m_a observers are sampled independently from a fixed query distribution ν. A layer is accepted only if every head satisfies η̂ + b ≤ ϵ_audit, where b = sqrt(ln(J·H_Q/δ) / (2m_a)) (Hoeffding plus a union bound over J monitored layers × H_Q heads). This holds even though the layer is picked adaptively. An optional stability streak (Jaccard-style set change ≤ θ) is a scheduling heuristic, not part of the guarantee. If no layer passes, full prefill runs.

flowchart TD
  A[Monitored layer] --> B{Doc mass / eligible heads valid?}
  B -- no --> R[Reset streak, continue unpruned]
  B -- yes --> C[Score chunks, greedy removal under caps]
  C --> D{Target + streak met?}
  D -- no --> R
  D -- yes --> E{Independent audit: η̂+b ≤ ϵ for all heads?}
  E -- no --> R
  E -- yes --> F[Gather states/KV on retained set, finish layers]
  R --> G{Window ended?}
  G -- yes --> H[Full-context prefill]

Evidence

The evidence is theoretical only. The bounds an implementer would use are:

• Pooled-to-head envelope. Any single proposal row's discarded mass is at most min{1, H_Q·W_p·q/ζ}, where q is the pooled discarded mass. The bound is sharp: a small pooled mass does not imply a small loss for a low-weight head. The H_Q·W_p/ζ factor is the real cost of pooling.
• Fixed-row output error. ‖o − o_R‖ ≤ 2Vη, with an attainable factor of 2. This is valid only for fixed queries, keys and values, not later layers.
• Counterexample (Prop. 6.1). Identical shallow attention can still produce different downstream outputs. No unconditional output guarantee follows from shallow observation.

Limits

• No experiments. The authors defer validation of the assumptions, the tightness of the bounds and end-to-end trade-offs.
• The audit bounds expected discarded mass under ν. It says nothing about the worst observer, any specific document position, or attention during decode.
• A deterministic last-token check does not qualify as an audit.
• Retuning the proposal using audit results invalidates the guarantee unless a fresh audit is drawn.
• If b exceeds the tolerance (m_a too small), no layer ever passes.
• Retaining m₀ chunks does not ensure a multi-hop reasoning chain survives.
• The entropy statistic is explicitly not a retrieval-head detector.
• The full-model stability bound needs additional Lipschitz and drift assumptions.
• The systems break-even (head unions vs. page allocation vs. KV-transfer payload) means gains from FlashAttention-2, PagedAttention and similar tools do not transfer automatically.
• The fetched text ends at Section 6.1. The perturbation theorem's constants and the break-even formula are not visible here.