Block-wise Prefill with Episodic Clustering

• The paper introduces a block-wise prefill strategy that segments data and clusters blocks into episodes, achieving up to 3.5x memory reduction in long-context LLMs.
• It employs unsupervised techniques like K-Means and dynamic programming to group data blocks into coherent episodes, ensuring retention of essential semantic context despite aggressive cache evictions.
• Adaptive layer-wise memory allocation based on sensitivity scores optimizes resource distribution across transformer layers, balancing efficiency and accuracy under strict memory constraints.

Block-wise prefill with episodic clustering refers to a family of techniques for managing memory, context, or structure in large-scale machine learning and data analysis tasks by segmenting data into blocks and grouping these blocks into coherent episodes (clusters) based on semantic, temporal, or application-specific criteria. These methods are particularly prominent in conversational LLM systems, few-shot learning, unsupervised segmentation, and episodic @@@@1@@@@, where both resource constraints and the preservation of long-term or structure-aware context are essential.

1. Conceptual Foundations

Block-wise prefill involves processing data (such as token sequences, matrix entries, or temporal series) in discrete blocks to control memory growth or to enable efficient intermediate computation. Episodic clustering, in this context, is the grouping of these blocks (or data segments) into semantically or statistically coherent episodes, often using unsupervised learning techniques (e.g., K-Means or dynamic programming-based partitioning).

Across domains, block-wise prefill mechanisms address the critical problem where naively accumulating all history—such as the full conversational context for LLMs, all timepoints for time series segmentation, or all past transitions for episodic control—rapidly exhausts available compute and memory budgets. Episodic clustering strengthens the effectiveness of these block-wise strategies by organizing blocks into meaningful units, allowing systems to retain, compress, or evict information at a granularity that is aligned with task-specific coherence objectives.

2. Block-wise Prefill in KV Cache Management

In long-context LLMs, especially those applied in multi-turn conversational question answering, the Key-Value (KV) cache grows linearly with history, resulting in substantial peak memory usage. The EpiCache framework (Kim et al., 22 Sep 2025) introduces block-wise prefill to manage this growth: instead of encoding the full conversational history before evicting non-essential tokens, the input is segmented into blocks of tokens (with size $M_{block}$). After each block is processed, only the most important KV entries (quantified through cross-attention scores with reference to a patched prompt derived from the dialogue) are retained, and the cache is trimmed back to the fixed size $M$. This ensures peak memory usage is bounded at $O(M + M_{block})$, as opposed to unbounded accumulation observed in post-prefill or full-context schemes.

This process is governed by importance scoring:

$$s_i = \max_{t \in [n+1, n+p]} \text{Attn}(x_t \rightarrow x_i)$$

using a patched prompt of length $p$, with $x_t$ denoting the prompt tokens and $x_i$ historical tokens.

The main advantage is the tight bounding of cache memory throughout token ingestion, which eliminates spikes and enables efficient inference for long conversations under strict resource constraints.

3. Episodic Clustering of Blocks

To prevent the loss of topic-relevant context inherent to aggressive cache compression, episodic clustering is employed (Kim et al., 22 Sep 2025). The conversational history is windowed into segments, each of which is embedded using a semantic encoder $f_{embed}$. These embeddings are then clustered (frequently via K-Means with k-means++ initialization) to extract episodic clusters:

$$\mathcal{C}(\{e_k\}) \rightarrow \{E_1, \ldots, E_E\}$$

Within each episode $E_e$, a medoid segment

$$S_{medoid}^{(e)} = \arg\max_{s \in E_e} \cos(e_s, C_e)$$

is chosen by maximizing the cosine similarity between segment embedding $e_s$ and cluster center $C_e$. This representative medoid supplies the patched prompt for block-wise prefill and scoring, ensuring that compression decisions retain the most salient historical content for the episode's topic.

This episodic block segmentation and clustering ensures that, when information is evicted, topic coherence is preserved and the context necessary for accurate downstream response is maintained.

4. Adaptive Layer-wise Memory Allocation

Block-wise prefill and episodic clustering are complemented by adaptive, layer-wise budget allocation. LLM transformer layers exhibit different sensitivities to cache entry pruning (Kim et al., 22 Sep 2025). Sensitivity for layer $l$ is quantified as:

$$s_l = 1 - \frac{1}{H \cdot N} \sum_{h=1}^{H} \sum_{i=1}^{N} \cos\left(K_{l,h,i}^{full}, K_{l,h,i}^{block}\right)$$

where $K^{full}$ and $K^{block}$ are Key state tensors under full-context and blockwise-evicted conditions, respectively; $H, N$ denote number of heads and sequence length.

Memory is then reallocated across layers according to:

$$M_{alloc}^{(l)} = \frac{s_l^\alpha}{\sum_{j=1}^L s_j^\alpha} (L \cdot M)$$

with exponent $\alpha$ tuning the distribution's sharpness. This ensures that layers most affected by token removal are allocated proportionally more cache, directly improving utility under memory constraints.

5. Episodic Clustering in Model-based Block Clustering and Time Series Segmentation

Model-based block clustering generalizes the block-wise episodic principle to data matrices and sequential data domains. The Bayesian collapsed latent block model (Wyse et al., 2010) samples both the row and column clusters and their number, with key elements:

• Block parameters (e.g., means, variances) are integrated out (“collapsed”) using conjugate priors, yielding a tractable, fixed-dimensional discrete posterior:

$$\pi(K, G, z, w \mid Y) \propto \pi(K) \pi(G) \frac{\Gamma(\alpha K)\prod_{k}\Gamma(n_k + \alpha)}{\Gamma(\alpha)^K \Gamma(n + \alpha K)} \frac{\Gamma(\beta G)\prod_g\Gamma(m_g + \beta)}{\Gamma(\beta)^G\Gamma(m + \beta G)} \prod_{k=1}^K\prod_{g=1}^G M_{kg}$$

where $M_{kg}$ is the integrated likelihood for block $(k, g)$.

• Sampling is performed using tailored MCMC moves—single row/column reassignments (Gibbs), block-wise reallocations, and split/combine proposals for changing the number of clusters.

This design enables the discovery and refinement of block structures (i.e., "prefilled" blocks) while handling their evolution across data episodes, naturally aligning with episodic clustering concepts.

In time-series segmentation, dynamic programming-based block clustering (Sinnathamby et al., 2021) imposes explicit constraints on block size and transition count, minimizing a cost summed over blocks such that each block is optimally coherent with respect to a physical or statistical model:

$$L_f(X, Y, W) = \sum_{c=0}^{C-1} \sum_{t:y_t=c} f(x_t, w_c)$$

Recurrence relations enforce both model coherence and temporal contiguity, resulting in segment (block) assignments that are meaningful within each modeled episode.

6. Applications in Episodic Control and Memory-Constrained Reinforcement Learning

In reinforcement learning, episodic control frameworks (Agostinelli et al., 2019) store compressed representations of past state-action transitions and employ these for rapid value estimation. To keep memory manageable, online clustering (notably, dynamic online k-means) is used: new transitions are merged into their nearest cluster, with a decay mechanism reducing cluster size if not recently updated. This maintains a compact and current episodic memory, ensuring the most relevant experience blocks are available for value approximation.

Dynamically adapting clusters to favor recent, high-value, or high-surprise transitions, episodic clustering enables efficient learning and mitigates catastrophic forgetting even under limited storage budgets. Block-wise prefill of the memory buffer, in this context, refers to only storing or updating clusters as new episodes are encountered, with old or less-relevant clusters decayed or replaced automatically.

7. Comparative Performance and Limitations

Techniques employing block-wise prefill and episodic clustering have demonstrated significant advantages:

Domain	Block/Episodic Mechanism	Demonstrated Gains
LLMs (LongConvQA)	EpiCache block-wise prefill + clustering	Up to 40% accuracy, 3.5x memory reduction (Kim et al., 22 Sep 2025)
Model clustering	Collapsed block MCMC	Homogeneous clusters, flexible K/G (Wyse et al., 2010)
Episodic RL	Dynamic online k-Means	Higher reward, small memory (Agostinelli et al., 2019)
Time series	DP segment clustering	Model-coherent, blockwise regimes (Sinnathamby et al., 2021)

A plausible implication is that while block-wise and episodic techniques improve memory management and context retention, overheads from clustering, episode management, and sensitivity estimation must be managed—though, for EpiCache, such overhead is reported at <5% of per-turn latency (Kim et al., 22 Sep 2025). Additionally, aggressive clustering or compression risks context aliasing or accuracy loss when not carefully balanced by sensitivity-aware mechanisms.

8. Broader Implications and Prospective Directions

Block-wise prefill with episodic clustering represents a structurally principled approach to scalable memory and context management in large-scale and long-horizon tasks. Its success in LLMs, reinforcement learning, and model-based segmentation highlights a common set of trade-offs: balancing resource use against the fidelity of contextual or episodic information. Adaptations, such as sensitivity-aware allocation and adaptive clustering, are critical in maintaining performance as model and data scale further. The underlying principle—aligning memory and context management with data structure and application episodes—suggests further applications in continual, multi-modal, or dynamically evolving data streams.