Architecture2024~103 minarXiv:2405.12981

Reducing Transformer Key-Value Cache Size with Cross-Layer Attention

score 9William Brandon, Mayank Mishra, Aniruddha Nrusimha +2↑ 33May 21, 2024cs.LG

🎯 Pitch

You can slash the KV cache memory of a large language model in half—without hurting quality—simply by making adjacent transformer layers share the same key and value projections. Cross-Layer Attention (CLA) achieves a consistent 2× memory reduction on top of Multi-Query Attention with less than 1% perplexity degradation, pushing the accuracy–memory frontier further than any head-sharing trick alone.

1. Executive Summary

This paper introduces Cross-Layer Attention (CLA), an architectural modification to the transformer that reduces the KV cache memory footprint by sharing key and value activations across adjacent layers—rather than computing separate KV projections at every layer, CLA designates a subset of layers to produce KV activations that neighboring layers reuse. Through pretraining experiments on 1B- and 3B-parameter models trained on SlimPajama, the authors systematically characterize CLA against Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) baselines, finding that CLA combined with MQA at a sharing factor of 2 (CLA2, where each KV projection is shared between pairs of consecutive layers) achieves a 2× reduction in KV cache size while incurring at most a "very modest (less than 1% change) degradation in perplexity"—for instance, an H128-MQA-CLA2 model requiring 5120 bytes per token matches the KV cache footprint of an H64-MQA baseline (also 5120 bytes) while improving validation perplexity by 0.21–0.31 points. The Pareto improvement holds across both parameter scales and survives learning rate tuning, establishing that sharing KV activations across layers advances the accuracy/memory frontier over MQA and GQA alone, though the benefits are most consistent when CLA2 is paired specifically with MQA rather than GQA, and sharing factors greater than 2 yield progressively worse tradeoffs.

2. Context and Motivation

The Core Problem: The KV Cache Is a Bottleneck at Scale

When a transformer-based language model generates text autoregressively—one token at a time—it must attend to all previously generated tokens. The naive approach of recomputing keys and values for every prior token at every generation step would be prohibitively expensive, so the standard optimization is to cache the key and value activations once they are computed and reuse them during subsequent steps. This key-value (KV) cache eliminates redundant computation, but its memory footprint scales linearly with three factors: sequence length, batch size, and the number of distinct key/value heads in the model.

The paper identifies this memory footprint as a bottleneck constraining real-world LLM deployment (Section 1):

"the memory overhead of KV cache storage can limit batch sizes when operating on long sequence lengths… and can require employing costly techniques like offloading when on-device memory is scarce"

This is not a minor inconvenience. Chowdhery et al. (2022) documented that KV cache memory can restrict throughput in serving systems. Sheng et al. (2023) showed that when on-device GPU memory is insufficient, practitioners resort to offloading KV cache entries to CPU or disk, which introduces substantial latency penalties. Beyond inference throughput, the paper points to an emerging use case that makes KV cache size even more critical: persistent KV caches. Systems like AttentionStore (Gao et al., 2024) and Google's Gemini context caching (Google, 2024) store KV caches across conversations to avoid recomputing attention over long shared-prefix contexts. In these scenarios, the storage footprint of the KV cache directly determines the cost of maintaining and retrieving conversational state.

The fundamental tension is this: as users demand longer context windows (100k+ tokens, for instance), and as serving systems aim for higher batch sizes to maximize throughput, the KV cache memory grows proportionally. GPUs have finite high-bandwidth memory (HBM), and the KV cache competes directly with model weights and activations for that scarce resource.

The Inefficiency: Redundant Key/Value Computation Across Layers

The paper opens by establishing that existing work has attacked the KV cache bottleneck along several dimensions—low-precision storage, eviction policies that discard less important KV entries, and sharing keys and values across query heads within a single layer. But the authors identify a dimension that prior work had not explored systematically: the fact that standard transformer architectures compute and store entirely separate KV projections for every single layer, even though adjacent layers might learn functionally similar representations. As the paper puts it in Section 1:

"we introduce a method for reducing the size of the KV cache along a dimension different than those explored in prior work: namely, reducing the number of unique layers in the KV cache."

The intuition is straightforward: if a transformer with 32 layers stores 32 separate KV caches, that's 32× the per-token KV memory of a single layer. But does every layer genuinely need unique KV projections? If layers can share key/value representations without meaningful degradation in model quality, the storage savings are immediate and proportional to the number of layers that become "KV consumers" rather than "KV producers."

This is the gap the paper identifies: prior work allocated the KV cache budget across query heads (via MQA and GQA) but never across layers. The observation that sharing could extend to the depth dimension is conceptually simple, but—critically—nobody had empirically characterized how this cross-layer sharing affects accuracy, what sharing patterns work best, or how it interacts with the existing MQA/GQA design space.

Why the Gap Has Practical Urgency

The memory overhead problem is not merely a theoretical concern; it is actively shaping architectural and deployment decisions. Consider the state of the field at the time of this paper's writing:

MQA and GQA are already widely adopted. The Llama series, Mistral, Gemma, and many other open-weight models use Grouped-Query Attention or Multi-Query Attention precisely because they reduce KV cache size compared to full Multi-Head Attention. These techniques are deployed because the memory savings matter in real serving scenarios.
The demand for long context is accelerating. Models like GPT-4-Turbo (128k tokens) and Claude 3 (200k tokens) push sequence lengths to new extremes. At 128k context with a typical 3B-parameter model, the KV cache can easily consume tens of gigabytes—comparable to or exceeding the model weights themselves.
Post-hoc compression has limits. Methods that quantize or evict KV entries after training (see the survey in Section 5.1) are complementary to architectural changes but operate within the information content already present in the cache. Architectural interventions like MQA and CLA reduce the memory footprint from the start of design, potentially preserving accuracy that post-hoc methods would sacrifice.

The paper therefore positions CLA not as a replacement for existing efficiency techniques but as an architectural primitive that pushes the Pareto frontier outward—enabling any given level of KV cache memory to achieve better accuracy than previously possible, or equivalently, enabling the same accuracy with a smaller memory budget.

Where Prior Approaches Fall Short

The paper's related work (Section 5) situates CLA against three categories of prior work on KV cache efficiency, and identifies specific limitations that motivate the need for a cross-layer approach.

1. Post-training compression techniques have inherent accuracy/cost tradeoffs.

Approaches like KVQuant (Hooper et al., 2024) and Coupled Quantization (Zhang et al., 2024) reduce KV cache memory by storing keys and values in low-precision formats (down to 1–2 bits). Cache eviction methods like H2O (Zhang et al., 2023), Scissorhands (Liu et al., 2023), and FastGen (Ge et al., 2024) selectively discard KV entries deemed unimportant. While effective, these are post-hoc interventions: they operate on representations produced by an already-trained model architecture. They are constrained by the information content in the original full-precision, full-size KV cache. There is an inherent accuracy penalty: you are discarding information (either precision or entire entries) that the model was trained to use. CLA, in contrast, is a training-time architectural change: the model learns from scratch to operate with a reduced KV cache structure, potentially developing representations that are more robust to the shared-KV constraint than what post-hoc compression can recover.

2. Attention replacements (SSMs, linear attention) abandon the softmax attention mechanism entirely.

Methods like Mamba (Gu and Dao, 2023), RWKV (Peng et al., 2024), and Gated Linear Attention (Yang et al., 2024) replace the quadratic softmax attention with state-space models or linear attention formulations that have constant memory complexity with respect to sequence length. These are promising but represent a complete departure from the transformer attention paradigm. They require retraining from scratch using architectures that may not yet match the quality of well-optimized softmax attention transformers at all scales and tasks. CLA preserves the core softmax attention computation—every layer still performs standard dot-product attention over keys and values—it only changes which layer produced those keys and values. This makes CLA a minimally invasive modification that is fully compatible with existing transformer training infrastructure, tensor parallelism, and the broader optimization ecosystem built around softmax attention.

3. Within-layer KV sharing (MQA, GQA) exhausts the "head" dimension but leaves the "layer" dimension untouched.

This is the most direct lineage of CLA. Shazeer (2019) proposed Multi-Query Attention: instead of having $n_{\text{query}}$ distinct key/value heads (one per query head, as in standard Multi-Head Attention), all query heads share a single key/value head. Ainslie et al. (2023b) generalized this to Grouped-Query Attention, where query heads are partitioned into groups, and each group shares a key/value head. The key metric is the number of distinct key/value heads per layer, denoted $n_{\text{group}}$ for GQA (with MQA being the special case $n_{\text{group}} = 1$ ).

These techniques reduce the per-layer KV cache size from $2 \cdot n_{\text{query}} \cdot d_{\text{head}}$ elements per token (MHA) to $2 \cdot n_{\text{group}} \cdot d_{\text{head}}$ (GQA). For a model with 16 query heads and $d_{\text{head}} = 128$ , switching from MHA to MQA (1 key/value head) reduces the per-layer KV cache by 16×. This is substantial. But it operates entirely within a single layer. Even with MQA, a 20-layer model still stores 20 separate key/value caches—one per layer. The total KV cache memory is:

$\text{KV cache size} = L \times 2 \cdot n_{\text{group}} \cdot d_{\text{head}} \cdot p$

where $L$ is the number of layers and $p$ is the precision (e.g., 2 bytes for 16-bit). MQA and GQA attack the $n_{\text{group}}$ factor. CLA attacks the $L$ factor by making $L$ in the equation above effectively mean "number of distinct key/value projections" rather than "number of layers." This is orthogonal: CLA can be combined with any setting of $n_{\text{group}}$ (MHA, GQA, or MQA), multiplying the memory reduction from each.

The paper explicitly draws this parallel in Section 2.2:

"Inspired by the success of MQA and GQA, which share key/value heads across query heads within a single layer, we propose also sharing key/value heads across layers."

The insight is that MQA/GQA proved that attention heads have redundancy across the query dimension—multiple query heads can productively attend over the same keys and values. CLA hypothesizes that attention layers have redundancy across the depth dimension—multiple layers can productively attend over the same keys and values as well.

What Makes the Cross-Layer Hypothesis Non-Obvious

It is worth emphasizing why this is not a trivial extension. In a standard transformer, each layer produces its own keys and values precisely because each layer's hidden representations differ. Layer 5's hidden states are not the same as layer 6's—they live in different representational spaces, refined by a full self-attention + MLP block. The question CLA poses is: if layer 6 uses the keys and values produced by layer 5's KV projection, can layer 6's query still extract useful information from them, despite the fact that those keys and values were computed from layer 5's (not layer 6's) hidden states?

Unlike MQA, where all query heads in a layer operate over the same hidden representation and therefore sharing KV projections is simply a capacity reduction, CLA asks layers to attend over KV representations that come from a different, earlier hidden representation. This is a more aggressive form of sharing because the input distribution to the attention computation changes: layer 6 queries attend over keys and values derived from layer 5's activations, not layer 6's. The model must learn to compensate for this mismatch. Whether this works effectively—and under what configurations—is an empirical question that the paper sets out to answer.

How the Paper Positions Itself

The paper frames CLA as a complementary extension of the MQA/GQA paradigm into the depth dimension, not as a replacement for any existing method. The explicit analogy (Section 2.2) is: just as GQA defines a family of attention configurations parameterized by $n_{\text{group}}$ (the number of key/value heads per layer), CLA defines a family parameterized by the sharing factor—the number of layers that share each KV projection. CLA2 (sharing between pairs of layers), CLA3 (sharing among triples), and so on, each represent different points in the accuracy/memory tradeoff space. Moreover, these parameterizations are fully composable with GQA's $n_{\text{group}}$ parameter, producing a two-dimensional design space (KV heads per layer × unique KV layers) where any point can be selected to match a target memory budget.

The paper's claim is modest and specific in scope: it does not argue that CLA is universally optimal, but rather that it expands the achievable Pareto frontier beyond what MQA/GQA alone can offer. As stated in Section 4:

"we find that MQA-CLA2 consistently achieves the lowest validation perplexity… for a given KV cache memory budget and model size"

The experimental design is correspondingly focused on matching baselines at equal memory rather than chasing absolute performance. A CLA model is evaluated against a non-CLA baseline that requires the same number of KV cache bytes per token—achieved by adjusting the baseline's head dimension or GQA factor to equalize memory. This is a fair comparison because it isolates the architectural innovation: given a fixed memory budget, does cross-layer sharing produce better representations than simply reducing the per-layer KV head count?

The Systems Perspective: What CLA Changes (and What It Doesn't)

Section 2.3 provides a pragmatic accounting of CLA's systems implications, which matters for understanding why this work is positioned as "practical engineering improvement" rather than "theoretical curiosity":

KV cache memory: significant reduction. The primary target. CLA shrinks the number of unique KV caches by approximately the sharing factor, yielding a proportional reduction in the total KV cache footprint.
Training memory: minor reduction. During training, intermediate KV tensors are materialized but are typically small compared to hidden states and MLP activations, especially for MQA/GQA models where $n_{\text{group}}$ is already small.
Parameters and FLOPs: slight reduction. Because CLA models have fewer distinct key/value projection weight matrices (they only exist in "KV-producing" layers), parameter count and forward/backward FLOPs decrease marginally. This is a secondary benefit, not the main motivation.
Core attention latency: no direct effect. This is a critical distinction from MQA/GQA. MQA reduces KV cache size and reduces memory bandwidth in the attention computation because fewer KV heads need to be read from memory per layer. CLA does not reduce per-layer memory bandwidth: "even shared KV cache layers must be separately re-read from main memory in each attention layer" (Section 2.3). CLA reduces storage, not per-step memory access. The latency benefits come indirectly: by reducing total KV cache memory, CLA enables larger batch sizes or longer sequences within the same memory budget, which improves overall throughput. But the per-token generation latency of a single sequence is not directly improved.

This distinction is important for understanding the paper's scope. CLA is not a general inference accelerator; it is specifically a memory capacity expansion technique. Its use case is scenarios where memory, not compute or memory bandwidth, is the binding constraint on deployment.

Summary of the Motivational Arc

The paper addresses a real and growing bottleneck: KV cache memory consumption limits the batch sizes and sequence lengths achievable in transformer inference. Existing mitigation strategies—post-hoc compression, architectural departures from attention, and within-layer KV sharing—leave one dimension unexplored: reducing the number of layers that require unique key/value caches. CLA fills this gap by proposing that adjacent layers can share KV activations without sacrificing accuracy, extending the MQA/GQA paradigm of controlled redundancy into the depth axis. The paper's empirical contribution is to rigorously characterize when and how this sharing works, establishing CLA2+MQA as a specific, actionable recipe that pushes the accuracy/memory Pareto frontier forward.

3. Technical Approach

3.1 Reader Orientation

This is a pretraining design-space exploration paper whose core idea is that key and value activations can be shared across adjacent transformer layers—not just across query heads within a single layer—yielding a proportional reduction in KV cache memory footprint with minimal accuracy degradation. The paper constructs a systematic empirical comparison between standard MQA/GQA baselines and novel "Cross-Layer Attention" (CLA) variants at matched KV cache memory budgets, trained from scratch at 1B and 3B parameter scales, to establish that CLA expands the achievable accuracy/memory Pareto frontier.

3.2 Big-Picture Architecture (Diagram in Words)

The CLA system is a modification to the standard transformer decoder architecture, composed of four interconnected components:

Base Transformer Architecture: A Llama-like decoder-only transformer with pre-normalization, SwiGLU activations, and rotary position embeddings—providing the backbone (layers, hidden size, FFN size) held constant across all compared variants.
Attention Mechanism (MHA / GQA / MQA): The within-layer attention design that determines how many distinct key/value projections exist per layer—ranging from full Multi-Head Attention (one KV head per query head) through Grouped-Query Attention (query heads partitioned into groups sharing KV heads) down to Multi-Query Attention (all query heads share a single KV head).
Cross-Layer Sharing Pattern: The new CLA component that designates a subset of layers as "KV producers" (which compute key and value projections from their hidden states) and the remaining layers as "KV consumers" (which reuse the KV activations from a designated earlier layer). This is parameterised by a sharing factor (2, 3, 4, etc.) and a sharing pattern (uniform pairing, keep-ends, dense-front, dense-back).
KV Cache Storage: The physical memory allocation that stores one set of key/value tensors per KV-producer layer—when CLA is active, the number of stored KV caches equals the number of KV producers, not the number of total layers.

Information flows as follows during autoregressive decoding: at each generation step, every layer computes its query projection from its own hidden state; KV-producer layers additionally compute fresh key and value projections (which are cached); KV-consumer layers instead read the cached keys and values of a designated earlier layer; all layers compute standard scaled dot-product attention using their own queries against the read keys and values; the attention output feeds into the MLP as usual; the total KV cache size is proportional to (number of KV producers) × (KV heads per producer) × (head dimension) × (precision).

3.3 Roadmap for the Deep Dive

First, the formal definition of CLA and its relationship to MQA/GQA (Section 2 material), establishing the shared mathematical notation and the two-dimensional design space (n_group, L_KV) that both techniques jointly parameterise.
Second, the precise mechanism of cross-layer KV sharing—which layers produce vs. consume KV activations, how the cache structure changes, and how the KV cache size formula is modified—because this is the core architectural intervention.
Third, the interaction between CLA and MQA/GQA—why the techniques are orthogonal and composable, and how the total KV cache size becomes the product of reductions along both the "heads per layer" and "layers" dimensions.
Fourth, the systems implications enumerated in Section 2.3, clarifying exactly what CLA changes (KV cache memory, parameter count, training FLOPs) and what it does not change (per-layer attention bandwidth, per-token generation latency), because understanding the scope is essential for knowing when CLA is applicable.
Fifth, the experimental configuration—model architectures, training recipe, shared hyperparameters, the CLA-specific design choices (which layers share, the KV-producer layer's layer norm parameters), and how models are compared at equal KV cache memory—since the empirical claims depend on fair comparison methodology.
Sixth, the learning rate tuning protocol and cross-validation approach that ensures baselines are not disadvantaged by suboptimal hyperparameters.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a pretraining design-space exploration paper whose core idea is that the standard transformer's per-layer key/value projections contain redundancy across the depth dimension, and that this redundancy can be eliminated by having layers share KV activations—analogous to how MQA/GQA eliminates redundancy across query heads within a layer.

The KV Cache Memory Budget: Formalising the Bottleneck

Before introducing CLA, the paper establishes the baseline KV cache size formula that all architectures (MHA, GQA, MQA, and CLA variants) will be measured against. The standard KV cache for a transformer stores, for every token in the sequence, the key and value activations for each distinct key/value head in each layer.

For a transformer with $L$ layers, $n_{\text{group}}$ distinct key/value heads per layer, head dimension $d_{\text{head}}$ , and storage precision of $p$ bytes per element, the KV cache size per token is:

$S_{\text{KV}} = L \times 2 \cdot n_{\text{group}} \cdot d_{\text{head}} \cdot p$

where $L$ is the number of transformer layers in the model, $n_{\text{group}}$ is the number of distinct key/value heads per layer (equal to $n_{\text{query}}$ for MHA, intermediate for GQA, and 1 for MQA), $d_{\text{head}}$ is the embedding dimension of each attention head, and $p$ is the number of bytes per scalar element (2 for 16-bit precision, as used throughout the paper's calculations).

What it computes: the total number of bytes that must be stored in GPU memory for the KV cache of a single token at a single generation step, summed over all layers. The factor of 2 accounts for storing both keys and values. This value is multiplied by sequence length and batch size to obtain the total KV cache memory for a generation batch.

Why this form: each layer independently computes and stores its own key and value projections. There is no sharing between layers, so the memory cost is additive across $L$ . The $n_{\text{group}}$ factor represents the per-layer design choice: MHA uses $n_{\text{group}} = n_{\text{query}}$ , MQA uses $n_{\text{group}} = 1$ , and GQA uses an intermediate value. This formula makes explicit that the KV cache size is linear in both the number of layers and the number of distinct KV heads per layer—and CLA's intervention is to effectively reduce the $L$ term.

Multi-Query Attention and Grouped-Query Attention: The Within-Layer Baseline

The paper builds directly on MQA and GQA, treating them as the established baselines against which CLA must demonstrate Pareto improvement. Understanding their mechanism is essential because CLA extends the same principle of controlled sharing to a new dimension.

In standard Multi-Head Attention (MHA), each of the $n_{\text{query}}$ query heads has its own dedicated key and value projection matrices $W^K_i, W^V_i$ , producing $n_{\text{query}}$ distinct sets of keys and values. The KV cache must store all of them, yielding $n_{\text{group}} = n_{\text{query}}$ .

Multi-Query Attention (MQA), introduced by Shazeer (2019), modifies this by having all query heads share a single key/value head. That is, there is only one $W^K$ and one $W^V$ projection per layer, producing one set of keys and one set of values. Every query head attends over this same KV pair. The KV cache stores only these two tensors per layer, giving $n_{\text{group}} = 1$ . The number of query heads $n_{\text{query}}$ can remain large (preserving the model's capacity to attend to different patterns), while the KV cache footprint collapses.

Grouped-Query Attention (GQA), introduced by Ainslie et al. (2023b), generalises MQA by partitioning the $n_{\text{query}}$ query heads into $n_{\text{group}}$ groups, where each group shares a single key/value head. MQA is the special case $n_{\text{group}} = 1$ , and MHA is the special case $n_{\text{group}} = n_{\text{query}}$ . GQA provides a continuous knob ( $n_{\text{group}}$ ) interpolating between the memory-efficient MQA extreme and the expressive MHA extreme.

The key empirical finding from prior work (which this paper takes as given) is that MQA and moderate GQA (e.g., GQA2, GQA4) incur only a small degradation in accuracy compared to MHA, while reducing KV cache size by a factor of $n_{\text{query}} / n_{\text{group}}$ . The paper's non-CLA baselines in Table 1 demonstrate this: at 1B scale with $d_{\text{head}} = 128$ , moving from MHA (KV bytes per token = 163,840, perplexity 13.15) to GQA4 (40,960 bytes, perplexity 13.36) to GQA2 (20,480 bytes, perplexity 13.52) to MQA (10,240 bytes, perplexity 13.54) shows the expected tradeoff—a 16× reduction in KV cache size for a perplexity increase of only 0.39 points.

CLA extends the MQA/GQA principle of KV sharing from the head dimension (within a single layer) to the depth dimension (across layers). The mechanism is defined in Section 2.2 and illustrated in Figures 1 and 2.

The core mechanism. In a standard transformer with $L$ layers, every layer $\ell$ applies its own learned key and value projections to its own hidden state $h_\ell$ :

$K_\ell = h_\ell W^K_\ell, \quad V_\ell = h_\ell W^V_\ell$

In a CLA transformer, only a subset of layers are designated as KV-producer layers—they retain their own $W^K$ and $W^V$ projection matrices and compute fresh keys and values from their own hidden states. The remaining layers are KV-consumer layers—they do not have $W^K$ or $W^V$ matrices and instead reuse the KV activations computed by a designated earlier (or, in principle, later) KV-producer layer.

Specifically, if layer $\ell$ is a KV consumer that reuses the KV cache of KV-producer layer $p$ (where $p < \ell$ ), then at both training and inference time, layer $\ell$ 's attention computation uses:

$K_\ell = K_p = h_p W^K_p, \quad V_\ell = V_p = h_p W^V_p$

Layer $\ell$ still computes its own queries from its own hidden state:

$Q_\ell = h_\ell W^Q_\ell$

And it still performs standard scaled dot-product attention:

$\text{Attention}(Q_\ell, K_p, V_p) = \text{softmax}\left(\frac{Q_\ell K_p^T}{\sqrt{d_{\text{head}}}}\right) V_p$

What happens physically. During autoregressive decoding, when a new token is generated, each layer's hidden state $h_\ell$ is computed sequentially (layer 0, then layer 1, etc.). When a KV-producer layer is reached, it computes $K_p$ and $V_p$ from $h_p$ and writes them to the KV cache. When a KV-consumer layer is reached, it reads the cached $K_p$ and $V_p$ from the designated earlier layer (which is already in memory, having been computed and cached earlier in the same forward pass) and uses them for attention. No new KV tensors are allocated or stored for the consumer layer.

The sharing factor and the number of KV caches. CLA is parameterised by a sharing factor $f$ , which is the number of consecutive layers that share a single KV cache. For a sharing factor of $f = 2$ (denoted CLA2), every pair of consecutive layers shares one KV projection: layer 0 produces KV for itself and layer 1; layer 2 produces KV for itself and layer 3; and so on. The number of distinct KV caches in the model becomes approximately $L / f$ (with slight adjustments if $f$ does not evenly divide $L$ ).

The modified KV cache size formula for CLA becomes:

$S_{\text{KV}}^{\text{CLA}} = L_{\text{KV}} \times 2 \cdot n_{\text{group}} \cdot d_{\text{head}} \cdot p$

where $L_{\text{KV}}$ is the number of KV-producer layers (the number of distinct KV caches), which equals approximately $\lceil L / f \rceil$ . This is the key equation: the original formula's $L$ factor is replaced by $L_{\text{KV}}$ , which is smaller by approximately the sharing factor $f$ .

Why this form works—the empirical hypothesis. The paper's central hypothesis is that adjacent layers in a transformer produce KV representations that are sufficiently similar, or that the query projections in consumer layers can learn to productively attend over KV representations from earlier layers, such that the information loss from sharing is small. This is not obvious a priori: the hidden states $h_p$ (from which $K_p$ and $V_p$ are computed) and $h_\ell$ (from which $Q_\ell$ is computed) live in different representational spaces—layer $p$ 's hidden states encode information after $p$ blocks of processing, while layer $\ell$ 's hidden states encode information after $\ell$ blocks. The query at layer $\ell$ is "looking for" different information than the query at layer $p$ , but it must do so using keys and values derived from an earlier, less-processed representation. The paper's experimental results (Section 3) answer whether this hypothesis holds in practice.

Why CLA is orthogonal to MQA/GQA. CLA operates on the $L_{\text{KV}}$ term (reducing the effective number of layers in the KV cache), while MQA/GQA operate on the $n_{\text{group}}$ term (reducing the number of distinct KV heads per layer). The total KV cache size with both techniques applied is:

$S_{\text{KV}}^{\text{CLA + GQA}} = L_{\text{KV}} \times 2 \cdot n_{\text{group}} \cdot d_{\text{head}} \cdot p$

Both $L_{\text{KV}}$ and $n_{\text{group}}$ can be independently varied, creating a two-dimensional design space where the total memory reduction is approximately $(L/L_{\text{KV}}) \times (n_{\text{query}} / n_{\text{group}})$ . The paper explores points in this space—for example, MQA-CLA2 with $d_{\text{head}} = 128$ has $n_{\text{group}} = 1$ (MQA) and $L_{\text{KV}} = 10$ (for $L = 20$ , CLA2), yielding a total KV cache of $10 \times 2 \times 1 \times 128 \times 2 = 5{,}120$ bytes per token—which matches the memory footprint of a plain MQA model with $d_{\text{head}} = 64$ (which has $L_{\text{KV}} = L = 20$ giving $20 \times 2 \times 1 \times 64 \times 2 = 5{,}120$ bytes).

The paper does not merely propose sharing KV activations; it investigates which specific layers should share. Section 3.2.1 describes four distinct sharing configurations, all evaluated against the uniform CLA2 baseline.

Uniform CLA2 (the default, and the best performer). With a 20-layer model, layers are paired: (0,1), (2,3), ..., (18,19). In each pair, the first layer (the even-indexed layer: 0, 2, 4, ...) is the KV producer—it has $W^K$ and $W^V$ projection matrices and writes to the KV cache. The second layer (the odd-indexed layer: 1, 3, 5, ...) is the KV consumer—it reads the KV cache of its paired predecessor. This yields $L_{\text{KV}} = 10$ KV caches out of $L = 20$ layers. The naming convention for the model is, for example, "H128-MQA-CLA2" meaning head dimension 128, MQA attention, CLA2 uniform sharing.

CLA2-KeepEnds (non-uniform, special treatment for boundaries). This configuration tests the hypothesis that the first and last layers might benefit from having their own dedicated KV caches, on the grounds that the input and output representations are qualitatively different from intermediate representations. Specifically, layer 0 retains its own KV cache (not shared with any other layer). Sharing then begins at layer 1: layers 1 and 2 share, layers 3 and 4 share, and so on. This also gives the final layer (layer 19) its own KV cache if the pairing works out, or at minimum ensures the last few layers are not forced to share. This configuration has $L_{\text{KV}} = 11$ —slightly more than uniform CLA2—because the "keep ends" constraint forces two layers to be unshared (layer 0 and one extra at the boundary). The model is denoted "H128-MQA-CLA2-KeepEnds."

CLA2-DenseFront (all sharing concentrated at the front). This configuration tests whether the model benefits from having all its KV-producing capacity concentrated in a contiguous block of early layers, which consumer layers then reference. It consists of 10 non-CLA layers (each with their own KV projections), followed by 9 CLA consumer layers that all reuse the KV activations of layer 9, followed by a final layer with its own KV cache. This yields $L_{\text{KV}} = 11$ (layers 0–9 produce, layer 10 reuses layer 9's KV, layers 11–19 reuse layer 9's KV, layer 19 has its own). The model is denoted "H128-MQA-CLA2-DenseFront."

CLA2-DenseBack (all sharing concentrated at the back). The mirror image: 2 non-CLA layers at the start, then a run of 10 CLA consumer layers all using the KV activations of layer 1, followed by 9 non-CLA layers. This yields $L_{\text{KV}} = 11$ as well. The model is denoted "H128-MQA-CLA2-DenseBack."

Why these patterns are tested. The uniform CLA2 pattern is the simplest and most memory-efficient (it minimises $L_{\text{KV}}$ to exactly $L/2$ when $L$ is even). But it might be suboptimal if, for example, consecutive layers have very different representational needs (the "keep ends" hypothesis) or if it is better to have a small number of "anchor" KV caches that many layers reference (the "dense" hypotheses). The experimental results (Table 1) show that all three alternative patterns perform worse than uniform CLA2 at the same or higher memory cost: CLA2-KeepEnds achieves perplexity 13.62 vs. 13.60 for uniform H128-MQA-CLA2 despite using more KV memory (5,632 vs. 5,120 bytes); DenseFront achieves 13.75; and DenseBack achieves 14.03—substantially worse. This provides empirical justification for the uniform pairing as the default CLA configuration.

The paper also explores CLA3 (sharing factor 3) and CLA4 (sharing factor 4), where groups of 3 or 4 consecutive layers share a single KV cache. These are described in Section 3.2.1 under "Ablation: MQA + CLA with Sharing Factor >2."

CLA3 with 20 layers. With a sharing factor of 3 and 20 layers, the layers are partitioned into groups of 3: (0,1,2), (3,4,5), ..., (18,19). However, 20 is not evenly divisible by 3: $20 / 3 = 6$ remainder 2. The paper's handling (Figure 2, rightmost diagram) is to have 6 full groups of 3 (covering layers 0–17) and then one remaining group covering layers 18–19 (only 2 layers), giving $L_{\text{KV}} = 7$ (one KV cache per group). The first layer in each group (layer 0, 3, 6, 9, 12, 15, 18) is the KV producer, and the remaining layers in the group are consumers. This yields a model denoted "H128-MQA-CLA3" with 3,584 bytes per token.

CLA4 with 20 layers. With a sharing factor of 4: $20 / 4 = 5$ groups of 4, exactly dividing 20. This yields $L_{\text{KV}} = 5$ KV caches, for a model "H128-MQA-CLA4" with 2,560 bytes per token.

Why higher sharing factors are tested. The paper's hypothesis is that sharing across more layers (CLA3, CLA4) should enable even greater KV cache reductions, but might encounter diminishing returns because a query in layer $\ell + 3$ attending over keys from layer $\ell$ perceives a much larger representational gap. The results (Table 1) confirm this: H128-MQA-CLA3 achieves perplexity 13.77 with 3,584 bytes/token, which is better than the plain H46-MQA baseline (13.96 at 3,680 bytes) but worse than the H128-MQA-CLA2 model extrapolated to similar memory. H128-MQA-CLA4 achieves 13.95 with 2,560 bytes/token—better than H32-MQA (14.37 at 2,560 bytes) but again worse than what CLA2 achieves at higher head dimensions. The paper's conclusion (Section 4): "we find that using sharing factors greater than 2 (CLA3 and above) achieves slightly worse accuracy/memory tradeoffs than using CLA2 and varying the head dimension." CLA2 is the sweet spot in the explored regime.

CLA Interaction with Separately-Learnable Layer-Norm Parameters

A subtle architectural detail mentioned in Section 3.1 is that CLA models use separately-learnable elementwise affine parameters for the layer-norm applied to KV projection blocks versus Q projection blocks in attention.

The mechanism. In a standard transformer, each attention block applies layer normalisation to its input hidden state before computing Q, K, and V projections. Typically, all three projections share the same layernorm output. In a CLA model, the KV-producing layer computes its layernorm on $h_p$ before projecting to $K_p$ and $V_p$ . The KV-consuming layer at position $\ell$ computes its own layernorm on $h_\ell$ before projecting to $Q_\ell$ , but uses the $K_p$ and $V_p$ from layer $p$ , which were computed from $h_p$ after a different layernorm transformation.

The paper's design choice is to give the KV-producer layer's layernorm and the consumer layer's layernorm separate learnable affine parameters (scale and shift). This means each layer has:

One set of layernorm parameters for its own Q projection (always present).
One set of layernorm parameters for its own K and V projections (present only in KV-producer layers, or present but unused in KV-consumer layers).

The paper specifies: "our CLA models use separately-learnable affine layer-norm parameters for the KV projection blocks and Q projection blocks in attention" (Section 3.1).

Why this matters. If the same layernorm parameters were used for the KV producer's input and the consumer's Q input, the model would be forced to use a single normalisation strategy for two different purposes—normalising the input to the KV projections (which serve downstream consumers) and normalising the input to the Q projections (which serve the current layer's attention). Separating them gives the model the flexibility to learn different normalisation statistics for each role. This is a minor detail but reflects the paper's attention to ensuring the CLA architecture has the necessary degrees of freedom to compensate for the cross-layer KV sharing.

Systems Implications: What CLA Changes and What It Does Not

Section 2.3 enumerates the systems-level consequences of adopting CLA, distinguishing between direct effects (due to architectural changes) and indirect effects (due to enabling larger batches/sequences within the same memory budget).

KV cache memory: primary benefit, direct reduction factor $\approx f$ . The number of distinct KV caches drops from $L$ to approximately $L / f$ , so the total KV cache memory per token drops by the same factor. For CLA2 with 20 layers, this is $20 \rightarrow 10$ , a 2× reduction. For CLA3, $20 \rightarrow 7$ , approximately 2.86×. The paper reports all KV cache sizes in "bytes per token at 16-bit precision," computed as $L_{\text{KV}} \times 2 \times n_{\text{group}} \times d_{\text{head}} \times 2$ (the final factor of 2 is for 2 bytes per 16-bit element).

Parameters and FLOPs: secondary, minor reduction. Because KV-consumer layers do not have $W^K$ or $W^V$ weight matrices, the total parameter count of the model slightly decreases. Each KV-producer layer has query, key, value, and output projection weight matrices ( $W^Q$ , $W^K$ , $W^V$ , $W^O$ ). Each KV-consumer layer has only $W^Q$ and $W^O$ . With $d_{\text{model}}$ as hidden size and $d_{\text{head}}$ as head dimension, each of $W^Q$ , $W^K$ , and $W^V$ has $d_{\text{model}} \times (n_{\text{group}} \times d_{\text{head}})$ parameters (for MQA, $n_{\text{group}} = 1$ , so this is $d_{\text{model}} \times d_{\text{head}}$ ). Removing these matrices for $(L - L_{\text{KV}})$ layers reduces parameter count and forward/backward FLOPs by a small fraction—not the primary motivation, but a minor added benefit.

Training memory footprint: minor reduction. During training, intermediate KV activation tensors are materialised (for gradient computation through attention). For MQA and GQA models, these tensors are already small relative to the hidden states and MLP activations because $n_{\text{group}}$ is small. CLA further reduces the count of these tensors by approximately $f$ , but the absolute memory savings during training are modest.

Core attention latency: no direct effect (critical distinction from MQA/GQA). MQA and GQA reduce not only KV cache storage but also the memory bandwidth required during each decoding step's attention computation—because there are fewer distinct KV heads to read from memory. CLA does not provide this benefit. As the paper states: "even shared KV cache layers must be separately re-read from main memory in each attention layer" (Section 2.3). In a CLA2 model, layer 1 still needs to read layer 0's cached keys and values from memory to compute attention—the read has to happen, even though the KV data is the same as what layer 0 already read. CLA reduces the storage footprint of the KV cache (how much GPU HBM is allocated for the cache) but not the per-step read bandwidth (how many bytes must be fetched per decoding step). The latency benefits are indirect: by freeing up memory, CLA enables larger batch sizes (improving throughput) or longer sequences (serving requests that would otherwise not fit), but the per-token generation latency of a single sequence is unchanged.

Model parallelism compatibility. CLA is "fully compatible with standard tensor parallelism techniques" for sharding model weights across accelerators, as stated in Section 2.3. In pipeline parallelism, layers that share a KV cache must either reside in the same pipeline stage (so the KV cache can be shared without communication) or the KV activations must be communicated between stages. This is a practical constraint but not a fundamental limitation—it affects how the model is partitioned across GPUs, not whether CLA works.

Persistence and reuse benefits. Because CLA reduces the total size of the KV cache, it proportionally reduces the cost of storing persistent KV caches for multi-turn conversations or long-prefix reuse. The paper explicitly connects to systems like AttentionStore and Gemini context caching: if KV caches are to be stored and retrieved over time, halving their size directly halves storage costs and retrieval bandwidth.

Why these distinctions matter for the paper's scope. CLA is positioned as a memory capacity expansion technique, not a general inference accelerator. Its use case is memory-constrained scenarios: long sequences that would otherwise cause out-of-memory errors, large batch sizes that would otherwise be infeasible, or persistent caching where storage costs dominate. Practitioners interested in reducing per-token latency should look to MQA/GQA, FlashAttention, or speculative decoding; practitioners hitting memory limits should consider CLA.

Experimental Configuration: Shared Training Recipe and Architectures

All experiments in Section 3 share a common training and architectural foundation, described in Section 3.1 and Table 2. This standardisation ensures that differences in perplexity are attributable to the attention architecture, not to confounding factors like different optimisers, data, or training duration.

Base architecture. All models use a Llama-like (Touvron et al., 2023) decoder-only transformer design with: pre-normalisation (layer norm applied before attention and MLP sublayers, not after), SwiGLU activations in the feed-forward network (a gated variant of the GELU activation, combining two linear projections with a sigmoid-gated linear unit), and rotary position embeddings (RoPE, which encode position information by rotating query and key vectors in a pairwise fashion, enabling better length generalisation). No dropout is used in any model.

Layer-norm parameters. All models use learnable elementwise affine parameters (scale and shift) for layer normalisation. As noted above, CLA models use separate affine parameters for the layernorm applied before KV projections versus before Q projections.

Query head count convention. Unless otherwise stated, the number of query heads $n_{\text{query}}$ is set such that $n_{\text{query}} \cdot d_{\text{head}}$ equals the hidden size $d_{\text{model}}$ . For example, a model with $d_{\text{model}} = 2048$ and $d_{\text{head}} = 128$ uses $n_{\text{query}} = 2048 / 128 = 16$ query heads. This constraint means that changing $d_{\text{head}}$ changes $n_{\text{query}}$ inversely—reducing head dimension increases the number of query heads, keeping the total query representation dimension constant.

Training data and tokenisation. All models are trained from scratch on the SlimPajama dataset (Soboleva et al., 2023), a 627B-token cleaned and deduplicated subset of RedPajama. Tokenisation uses the GPT-NeoX tokenizer (Black et al., 2022), which employs Byte-Pair Encoding (BPE). The training data order is held consistent across all 1B-scale experiments, eliminating data-order variance as a confound.

Optimisation hyperparameters. The paper uses AdamW (Loshchilov and Hutter, 2019) with gradient clipping. The specific settings are:

$\beta_1 = 0.9$ , $\beta_2 = 0.95$ (Adam moment decay rates)
Weight decay factor of 0.1 (the decoupled weight decay regularisation strength)
Gradient clipping norm of 1.0 (gradients are rescaled if their L2 norm exceeds 1.0)
Linear learning rate warmup for the first 5% of training steps
Cosine learning rate schedule decaying to 10% of the peak learning rate over the remaining 95% of training

Batch configuration. Sequence length is 2,048 tokens. Batch size is 2,048 sequences. This gives approximately 4 million tokens per training step ( $2048 \times 2048 = 4{,}194{,}304$ ).

Initialisation. All linear layer weights are initialised from a normal distribution with mean zero and standard deviation 0.01275—a small value consistent with standard transformer initialisation practices to keep initial activations well-conditioned.

Hardware and precision. All experiments run on NVIDIA H100 GPUs using PyTorch (Paszke et al., 2019; Ansel et al., 2024). Training uses mixed precision in BF16 (Brain Floating Point 16, which has the same exponent range as FP32 but fewer mantissa bits than FP16, providing better dynamic range for training stability) with gradient all-reduce and gradient accumulation performed in FP32.

Model scale hyperparameters (Table 2).

Parameter	1B Models	3B Models
Hidden size $d_{\text{model}}$	2048	3072
FFN size	5472	8192
Number of layers $L$	20	32
Sequence length	2048	2048
Training tokens	30 billion	100 billion

All models at each scale share these hyperparameters, meaning the comparison is controlled: a 1B CLA model has the same hidden size, FFN size, number of total layers, and training token budget as a 1B non-CLA baseline. The only difference is the attention mechanism (CLA vs. not, sharing factor, head dimension, MQA vs. GQA factor).

Fair Comparison Methodology: Matching KV Cache Memory Budgets

The central empirical strategy of the paper is to compare CLA models against non-CLA baselines at equal KV cache memory footprint. This requires careful construction of comparison pairs.

Design space exploration logic (Section 3.2.1). The paper first trains a "ladder" of non-CLA baselines spanning different KV cache sizes by varying $d_{\text{head}}$ and the GQA factor. For example, at 1B scale with MQA:

H128-MQA: $d_{\text{head}} = 128$ , $L_{\text{KV}} = 20$ , $n_{\text{group}} = 1$ → KV bytes = $20 \times 2 \times 1 \times 128 \times 2 = 10{,}240$
H64-MQA: $d_{\text{head}} = 64$ , $L_{\text{KV}} = 20$ , $n_{\text{group}} = 1$ → KV bytes = $20 \times 2 \times 1 \times 64 \times 2 = 5{,}120$
H46-MQA: $d_{\text{head}} = 46$ , $L_{\text{KV}} = 20$ , $n_{\text{group}} = 1$ → KV bytes = $20 \times 2 \times 1 \times 46 \times 2 = 3{,}680$
H32-MQA: $d_{\text{head}} = 32$ , $L_{\text{KV}} = 20$ , $n_{\text{group}} = 1$ → KV bytes = $20 \times 2 \times 1 \times 32 \times 2 = 2{,}560$

Then CLA models are trained with matching KV cache sizes by adjusting $d_{\text{head}}$ upward to compensate for the reduced $L_{\text{KV}}$ . For CLA2 with MQA and $L_{\text{KV}} = 10$ :

H128-MQA-CLA2: $d_{\text{head}} = 128$ , $L_{\text{KV}} = 10$ → KV bytes = $10 \times 2 \times 1 \times 128 \times 2 = 5{,}120$ (matches H64-MQA)
H90-MQA-CLA2: $d_{\text{head}} = 90$ , $L_{\text{KV}} = 10$ → KV bytes = $10 \times 2 \times 1 \times 90 \times 2 = 3{,}600$ (matches H46-MQA approximately)
H64-MQA-CLA2: $d_{\text{head}} = 64$ , $L_{\text{KV}} = 10$ → KV bytes = $10 \times 2 \times 1 \times 64 \times 2 = 2{,}560$ (matches H32-MQA)
H256-MQA-CLA2: $d_{\text{head}} = 256$ , $L_{\text{KV}} = 10$ → KV bytes = $10 \times 2 \times 1 \times 256 \times 2 = 10{,}240$ (matches H128-MQA)
H512-MQA-CLA2: $d_{\text{head}} = 512$ , $L_{\text{KV}} = 10$ → KV bytes = $10 \times 2 \times 1 \times 512 \times 2 = 20{,}480$ (matches H128-GQA2)

Why this comparison design is fair. A CLA2 model with $d_{\text{head}} = 128$ and a plain MQA model with $d_{\text{head}} = 64$ use exactly the same amount of KV cache memory per token. If the CLA2 model achieves better perplexity, it means that cross-layer sharing with larger heads is a more efficient use of memory than per-layer unique caches with smaller heads. The comparison isolates the architectural intervention (CLA vs. non-CLA) from the memory budget—any perplexity difference is attributable to how the memory is structured, not how much memory is used.

Why head dimension is adjusted rather than other parameters. The paper could have matched memory budgets by adjusting $n_{\text{group}}$ (e.g., CLA2 with GQA vs. plain MQA), and indeed does this for some comparisons (e.g., H256-GQA4-CLA2 matching the memory of H128-GQA4). But varying $d_{\text{head}}$ is the most natural knob because it directly trades off per-head representational capacity against KV cache size—halving $d_{\text{head}}$ halves the KV cache (all else equal). CLA2 also halves the KV cache, so CLA2 with $d_{\text{head}} = X$ naturally matches plain with $d_{\text{head}} = X/2$ .

Validation and Evaluation Metrics

The paper uses two primary accuracy metrics and two types of memory metrics.

Validation perplexity. The primary metric for comparing model quality during the design space exploration. After training on 30B tokens (1B scale) or 100B tokens (3B scale), the model's perplexity is computed on a held-out validation set of approximately 4 million tokens drawn from the SlimPajama corpus. Perplexity is the exponentiated average negative log-likelihood:

$\text{PPL} = \exp\left(-\frac{1}{T} \sum_{t=1}^{T} \log p_\theta(x_t | x_{<t})\right)$

where $T$ is the number of validation tokens, $x_t$ is the $t$ -th token, and $p_\theta(x_t | x_{<t})$ is the model's predicted probability for the correct next token given the preceding context. Lower perplexity is better, and it directly measures the model's compression efficiency (how surprised it is by the held-out data).

Wikitext perplexity. A secondary metric used in the learning-rate-tuned model comparisons (Tables 3, 5, 7). Wikitext (Merity et al., 2016) is a standard language modelling benchmark consisting of Wikipedia articles. It provides an out-of-distribution (relative to SlimPajama) measure of generalisation quality.

Downstream benchmarks (Tables 4, 6, 8). The paper evaluates learning-rate-tuned models on seven standard benchmarks using EleutherAI's LM Eval Harness: HellaSwag (commonsense reasoning), PIQA (physical commonsense reasoning), WinoGrande (pronoun resolution), SciQ (science question answering), OpenBookQA (multi-step reasoning with science facts), BoolQ (boolean question answering), and ARC-Easy (grade-school science multiple choice). These provide task-level signal beyond perplexity.

KV cache memory (bytes per token at 16-bit). The primary memory metric, computed as $L_{\text{KV}} \times 2 \times n_{\text{group}} \times d_{\text{head}} \times 2$ . This is the per-token storage requirement in bytes for the KV cache at 16-bit (2-byte) precision. It is the independent variable on the x-axis of Pareto frontier plots (Figure 3), with validation perplexity on the y-axis.

Learning Rate Tuning Protocol

Because "the relative performance of different model architectures can change depending on the learning rates at which they are evaluated" (Section 3.2.2), the paper conducts a systematic learning rate sweep on key models to verify that CLA benefits are not an artifact of suboptimal learning rates for baselines.

Procedure. For each selected model configuration, the learning rate is swept upward from an initial value of $\text{LR} = 3 \times 10^{-4}$ in multiplicative increments of 1.5×. The sweep continues until validation perplexity stops improving (indicating the optimal learning rate has been passed). The learning rate that achieves the lowest validation perplexity is treated as an approximation of the model's optimal learning rate.

Models selected for tuning at 1B scale. Three models: H128-MQA (the $d_{\text{head}} = 128$ MQA baseline), H64-MQA (the $d_{\text{head}} = 64$ MQA baseline that matches CLA2 memory), and H128-MQA-CLA2 (the CLA model). These three represent the key comparison: CLA2 model vs. equal-head-dimension baseline (2× memory) and equal-memory baseline (smaller heads).

Results of the 1B sweep (Table 3).

H128-MQA: optimal LR = $1.5 \times 10^{-3}$ , validation perplexity = 12.39
H128-MQA-CLA2: optimal LR = $2.25 \times 10^{-3}$ , validation perplexity = 12.43
H64-MQA: optimal LR = $2.25 \times 10^{-3}$ , validation perplexity = 12.74

The qualitative result from the initial design space exploration (all models at $\text{LR} = 3 \times 10^{-4}$ ) holds at optimal learning rates: CLA2 incurs only a 0.04 point perplexity degradation vs. the 2×-memory H128-MQA baseline, and achieves a 0.31 point improvement over the equal-memory H64-MQA baseline.

Notable finding: CLA benefits from higher learning rates. The optimal learning rate for H128-MQA-CLA2 ( $2.25 \times 10^{-3}$ ) is higher than for H128-MQA ( $1.5 \times 10^{-3}$ ). The paper speculates in Appendix A that CLA models may benefit from higher learning rates—a pattern partially replicated at 3B scale where H128-MQA-CLA2's optimal LR is $2.25 \times 10^{-3}$ vs. H128-MQA's $6.75 \times 10^{-4}$ . This is a practical consideration: if practitioners adopt CLA, they should not assume the same learning rate as the non-CLA baseline is optimal.

3B-scale tuning (Section 3.3). The same protocol is applied at 3B scale, with two sets of experiments:

First set: H128-MQA (LR = $6.75 \times 10^{-4}$ , PPL = 9.52), H128-MQA-CLA2 (LR = $2.25 \times 10^{-3}$ , PPL = 9.34), H64-MQA (LR = $1.0 \times 10^{-3}$ , PPL = 9.48). An unexpected result: CLA2 actually outperforms the 2×-memory baseline (9.34 vs. 9.52), and the smaller-head H64-MQA (9.48) outperforms H128-MQA (9.52)—suggesting that at 3B scale, the $d_{\text{head}} = 128$ baseline may be undertrained or suboptimally configured.
Second set (on different hardware with different data order): H64-MQA (LR = $1.0 \times 10^{-3}$ , Wikitext PPL = 12.94), H64-MQA-CLA2 (LR = $1.0 \times 10^{-3}$ , PPL = 12.99), H32-MQA (LR = $1.0 \times 10^{-3}$ , PPL = 13.34). This replicates the 1B-scale pattern: CLA2 incurs 0.05 point degradation vs. 2×-memory baseline and 0.35 point improvement vs. equal-memory baseline.

Summary of Design Choices and Their Justifications

Llama-like architecture (pre-norm, SwiGLU, RoPE): matches contemporary best practices in language model design, ensuring results are representative of modern LLM configurations rather than legacy architectures.
SlimPajama training data: a large, publicly available corpus that enables reproducible pretraining experiments without proprietary data dependencies.
Training from scratch (not fine-tuning or conversion from MHA checkpoints): isolates the effect of the CLA architecture itself, avoiding confounding factors from checkpoint conversion or knowledge distillation, though it incurs higher computational cost.
KV cache memory as the equalisation metric: ensures fair comparison by matching the resource that CLA is designed to reduce, rather than comparing at equal head dimension or equal parameter count.
Head dimension as the primary adjustment knob: naturally trades off representational capacity against memory, and is continuous enough to produce a smooth Pareto frontier.
Uniform CLA2 as the default sharing pattern: simplest to implement and most memory-efficient, and empirically validated as the best performer among tested patterns.
Separate layernorm parameters for Q vs. KV projections: gives the model flexibility to learn different normalisation statistics for the two different roles, which is necessary when the KV projections serve downstream consumers while Q projections serve the current layer.
Learning rate tuning on key comparisons: addresses the threat that CLA's apparent benefits might be an artifact of using a single learning rate that happens to favour CLA, strengthening the validity of the Pareto improvement claims.
Two-parameter-scale replication (1B and 3B): provides evidence that CLA benefits are not specific to a single model size, though the unexpected 3B results (H64-MQA outperforming H128-MQA) suggest more complex scaling behaviour that warrants further study.

4. Key Insights and Innovations

The paper's central conceptual move is to identify the layer dimension as an independent, previously unexploited axis for KV cache reduction, orthogonal to the well-established head dimension (MQA/GQA). Prior work treated the KV cache size formula S_{\text{KV}} = L \times 2 \cdot n_{\text{group}} \cdot d_{\text{head}} \cdot p as having two tunable knobs: n_{\text{group}} (how many distinct KV heads per layer, controlled by MQA/GQA) and d_{\text{head}} (head dimension). The number of layers L was taken as a fixed architectural constant—every layer simply got its own KV cache because that's how transformers had always been built. CLA reframes this: L becomes L_{\text{KV}}, the number of KV-producing layers, which can be substantially smaller than the total number of layers. The formula becomes a two-dimensional design space where the total memory reduction is the product of reductions along both axes: (L/L_{\text{KV}}) \times (n_{\text{query}} / n_{\text{group}}).

This is not merely an incremental parameter tweak. It changes the mental model of what a transformer layer requires. Before CLA, the implicit assumption was that each layer's keys and values must be derived from that layer's own hidden states because each layer occupies a different representational space. CLA demonstrates that this assumption is overly conservative: layers can productively attend over KV representations from earlier, less-processed hidden states. The fact that this works—and works with minimal accuracy degradation at CLA2—means that the field's default of per-layer KV projections is a point of significant redundancy, not a fundamental requirement. The paper does not discover this redundancy post-hoc through analysis of trained models; it proposes the architectural change a priori and validates it through controlled pretraining, establishing CLA as a design primitive rather than a compression afterthought.

The significance extends beyond the specific CLA2 recipe. By parameterising the layer-sharing dimension (sharing factor f, sharing pattern), the paper opens a new axis for architecture search that is fully composable with existing head-sharing techniques. A model designer now has a grid of (n_{\text{group}}, L_{\text{KV}}) configurations to explore, where previously they had only a line of n_{\text{group}} options. This expansion of the design space is a fundamental contribution even if specific configurations (CLA3, CLA4) prove suboptimal—the framework itself is the insight.

Innovation 2: The Accuracy/Memory Pareto Frontier Reframing—Why "Equal Memory" Comparison Matters More Than "Equal Architecture"

The paper's evaluation methodology is itself an intellectual contribution: comparing architectures at equal KV cache memory footprint rather than at equal head dimension, equal parameter count, or equal FLOPs. This seems obvious in retrospect, but it represents a shift from how attention architecture research was often conducted. Prior MQA and GQA papers tended to compare at equal model size (same d_{\text{model}}, same L, same d_{\text{head}}) and reported perplexity degradation as the cost of memory savings—framing the tradeoff as "you sacrifice some accuracy to save memory." The CLA paper inverts this: it asks, given a fixed memory budget, what is the best architecture?

This reframing surfaces the key Pareto improvement that would be invisible under an equal-head-dimension comparison. An H128-MQA-CLA2 model (5,120 bytes/token, perplexity 13.60) compared to an H128-MQA model (10,240 bytes/token, perplexity 13.54) would look like a degradation—you're losing 0.06 perplexity points. But that's comparing at equal head dimension, not equal memory. The fair comparison is H128-MQA-CLA2 (5,120 bytes, 13.60) vs. H64-MQA (5,120 bytes, 13.81): CLA2 wins by 0.21 points at the same memory. The learning-rate-tuned version widens this to 0.31 points (Table 3). Under the conventional comparison framework, this gain is invisible; under the memory-matched framework, it's the headline result.

This methodological choice matters because it mirrors the real-world constraint that practitioners face: GPU HBM is a fixed capacity. The decision is not "should I trade 0.06 perplexity for 2× memory savings?" but rather "given 5,120 bytes/token of KV cache budget, should I use a plain MQA model with small heads or an MQA-CLA2 model with larger heads?" The paper demonstrates that the latter consistently dominates, which is the definition of a Pareto improvement. By adopting this comparison framework, the paper implicitly argues that the field should evaluate KV cache efficiency techniques on a memory-equivalent basis, not an architecture-equivalent basis—a standard that has since been adopted by subsequent work but was not the norm when CLA was proposed.

A non-obvious empirical finding with implications for architecture design is that CLA is most effective when combined specifically with MQA, not with GQA or MHA. Table 1 shows that GQA4-CLA2 with d_{\text{head}} = 128 achieves perplexity 13.48 with 20,480 bytes/token, which is slightly worse than the plain GQA4 baseline with d_{\text{head}} = 128 (13.36 at 40,960 bytes) when accounting for the 2× memory difference—the comparison point would be H64-GQA4, which is not in the table but would be at 20,480 bytes. More tellingly, H256-GQA4-CLA2 (40,960 bytes, 13.38) essentially matches plain H128-GQA4 (40,960 bytes, 13.36) rather than improving on it. GQA2-CLA2 (10,240 bytes, 13.59) is essentially tied with H128-MQA-CLA2 at the same memory (10,240 bytes, 13.51 for H256-MQA-CLA2; Table 1 shows 13.59 for GQA2-CLA2 vs. 13.51 for H256-MQA-CLA2 at 10,240 bytes—a slight disadvantage for GQA).

The pattern suggests that the benefits of cross-layer sharing are most pronounced when the per-layer KV capacity is already minimised (MQA, one KV head per layer) and the model compensates by keeping head dimension larger. When per-layer KV capacity is higher (GQA2, GQA4), the relative gain from layer sharing diminishes—possibly because GQA already provides enough per-layer expressive range that the bottleneck is not the number of distinct KV representations but rather the head dimension or query-side capacity. This is a practically actionable finding: if a practitioner is already committed to MQA for its memory efficiency, CLA2 is an unambiguously good addition. If they are using GQA, the case for CLA is weaker and requires more careful evaluation at the specific memory target.

This result also suggests a deeper principle about architectural redundancy: redundancy across heads (what GQA removes) and redundancy across layers (what CLA removes) may be partially substitutable. Eliminating both simultaneously yields diminishing returns, while concentrating the reduction on one dimension (heads, via MQA) and preserving capacity on the other (larger d_{\text{head}} enabled by CLA's layer reduction) produces the best tradeoff. This is not a theoretical claim the paper makes explicitly, but it emerges from the data and represents a testable hypothesis about transformer attention design that future work could exploit.

The paper's ablation of sharing factors greater than 2 yields a finding that is both practically useful and theoretically informative: CLA2 dominates CLA3 and CLA4 at matched memory budgets. H128-MQA-CLA3 achieves perplexity 13.77 with 3,584 bytes/token, while H90-MQA-CLA2 achieves 13.73 with 3,600 bytes/token—essentially the same memory but CLA2 wins on perplexity. H128-MQA-CLA4 achieves 13.95 with 2,560 bytes/token, while H64-MQA-CLA2 achieves 13.89 at the same 2,560 bytes—again CLA2 wins.

This is significant because it runs counter to a naive extrapolation: if sharing across 2 layers works well, sharing across 3 or 4 should work even better for memory reduction. The data says no—there is a non-monotonic relationship between sharing factor and efficiency at equal memory. The paper's explanation is implicit but interpretable from the architecture: when the sharing factor grows, the representational gap between the KV producer's hidden state and the consumer's hidden state widens. A query at layer \ell + 3 attending over keys from layer \ell faces a much larger mismatch than a query at layer \ell + 1 attending over layer \ell's keys. The model can compensate by learning query projections that bridge this gap, but there appears to be a limit—somewhere between sharing factor 2 and 3, the compensation cost exceeds the benefit of using a larger head dimension enabled by the memory savings.

This finding provides a concrete design guideline (use CLA2, not higher sharing factors) and also constrains future research: efforts to push cross-layer sharing further likely need additional mechanisms (e.g., learned transformations between producer and consumer KV representations, or adaptive sharing patterns) rather than simply increasing the sharing factor. The CLA2 sweet spot also aligns with an intuitive architectural prior: pairs of consecutive layers have the most similar representations, so pair-wise sharing extracts the "easiest" redundancy first, and there is relatively little additional redundancy to extract from wider groupings.

Innovation 5: Learning Rate Sensitivity as a Diagnostic Signal for Architectural Change

A subtle but methodologically important finding is that CLA models systematically benefit from higher learning rates than their non-CLA counterparts. At 1B scale, H128-MQA-CLA2's optimal learning rate is 2.25 \times 10^{-3} versus 1.5 \times 10^{-3} for H128-MQA—a 1.5× difference. At 3B scale, the gap is even larger: 2.25 \times 10^{-3} for CLA2 versus 6.75 \times 10^{-4} for the non-CLA baseline—a 3.3× difference (Tables 3 and 5). This pattern is not fully explained in the paper but has important implications.

The most plausible interpretation is that CLA models have fewer distinct key/value projection weight matrices to optimise (only in KV-producer layers), meaning the optimisation landscape has fewer parameters in the attention mechanism and a correspondingly different conditioning structure. Higher learning rates may be tolerable—or necessary—because the reduced parameter count changes the scale of the gradients flowing through the attention sublayers. This is consistent with the observation that CLA2 slightly reduces total parameter count and training FLOPs (Section 2.3), but the learning rate effect is larger than what would be expected from a small parameter reduction alone.

The practical significance is that naively applying a baseline's tuned learning rate to a CLA model will underestimate CLA's performance. If the paper had used LR = 1.5 \times 10^{-3} for H128-MQA-CLA2 (the optimal LR for the H128-MQA baseline) instead of tuning it, the CLA model's perplexity would have been worse than reported, potentially erasing the Pareto improvement. This is a methodological trap that the paper's learning rate sweeps deliberately avoid, and it serves as a cautionary example for architecture research more broadly: when you change the architecture, you change the optimisation dynamics, and hyperparameters must be re-tuned for fair comparison. The paper demonstrates this discipline and surfaces the learning rate shift as a finding in its own right—not just a nuisance variable to control for.

5. Experimental Analysis

Evaluation Methodology

Dataset. All models are trained on the SlimPajama dataset (Soboleva et al., 2023), a 627B-token cleaned and deduplicated version of RedPajama, tokenized with the GPT-NeoX Byte-Pair Encoding tokenizer (Black et al., 2022). Validation perplexity is computed on a held-out set of approximately 4 million tokens drawn from the same SlimPajama corpus. For the learning-rate-tuned models, additional evaluation is performed on Wikitext (Merity et al., 2016) perplexity and seven downstream benchmarks (HellaSwag, PIQA, WinoGrande, SciQ, OpenBookQA, BoolQ, ARC-Easy) using EleutherAI's LM Eval Harness (Gao et al., 2023).
Base model(s). All experiments use a Llama-like (Touvron et al., 2023) decoder-only transformer architecture with pre-normalization, SwiGLU activations (Shazeer, 2020; Ramachandran et al., 2017), and rotary position embeddings (Su et al., 2023). At 1B scale, models have hidden size 2048, FFN size 5472, and 20 layers. At 3B scale, models have hidden size 3072, FFN size 8192, and 32 layers. The architecture is chosen to represent contemporary LLM design practice, and training from scratch avoids confounding factors from checkpoint conversion or knowledge distillation approaches used in some prior GQA work.
Metrics. The primary metric is validation perplexity on a held-out SlimPajama subset, computed as the exponentiated average negative log-likelihood across validation tokens. Secondary metrics include Wikitext perplexity and downstream benchmark accuracy (percentage correct on multiple-choice tasks). The resource metric is KV cache memory, computed as bytes per token at 16-bit precision: L_KV × 2 × n_group × d_head × 2, where L_KV is the number of KV-producer layers, n_group is the number of distinct key/value heads per layer, d_head is the head dimension, and the final factor of 2 accounts for 2 bytes per 16-bit scalar.
Baselines. The paper trains seven non-CLA baselines spanning the MHA–GQA–MQA spectrum at varying head dimensions (Table 1):
- H128-MHA: Multi-Head Attention with d_head = 128, 16 query heads, 16 KV heads, 20 KV layers — the largest KV cache footprint (163,840 bytes/token).
- H128-GQA4: Grouped-Query Attention with 4 groups, d_head = 128, 16 query heads — 40,960 bytes/token.
- H128-GQA2: GQA with 2 groups, d_head = 128, 16 query heads — 20,480 bytes/token.
- H128-MQA: Multi-Query Attention with d_head = 128, 16 query heads, 1 KV head — 10,240 bytes/token.
- H64-MQA: MQA with d_head = 64, 32 query heads — 5,120 bytes/token.
- H46-MQA: MQA with d_head = 46, 45 query heads — 3,680 bytes/token.
- H32-MQA: MQA with d_head = 32, 64 query heads — 2,560 bytes/token.
Generation budget / compute accounting. All models at a given scale are trained on the same number of tokens (30B for 1B, 100B for 3B) with the same batch size (2048 sequences of 2048 tokens, ~4M tokens per step). CLA models require slightly fewer FLOPs during training due to fewer key/value projection matrices, but the paper does not adjust training tokens to equalize FLOPs — the comparison is at equal token budget, which slightly favors CLA (it gets effectively the same number of optimizer steps at marginally lower computational cost). KV cache memory is equalized across comparison pairs by adjusting head dimension and/or the GQA grouping factor.
Cross-validation / statistical protocol. There is no formal cross-validation or statistical significance testing reported. The learning rate tuning protocol (Section 3.2.2) serves as the primary robustness check: for key model configurations, learning rates are swept upward in multiplicative increments of 1.5× from LR = 3 × 10^{-4} until validation perplexity stops improving, and the best learning rate for each architecture is used in the final comparison. This ensures that baseline models are not disadvantaged by suboptimal hyperparameters. The 3B-scale replication on different hardware with different data order (Section 3.3, second set) serves as a partial robustness check against infrastructure-specific variance.

Main Quantitative Results

1B-Scale Design Space Exploration (Table 1, Figure 3)

The 1B-scale design space exploration is the paper's primary empirical contribution, training 24 models (7 baselines + 17 CLA variants) on 30B tokens each at a conservative learning rate of LR = 3 × 10^{-4}.

Headline result: CLA2+MQA advances the Pareto frontier. Figure 3 plots validation perplexity against KV cache bytes per token for all models. The key comparison pairs are CLA2+MQA models matched in memory to plain MQA baselines by halving the head dimension of the baseline relative to the CLA variant:

H128-MQA-CLA2 (5,120 bytes/token, perplexity 13.60) vs. H64-MQA (5,120 bytes/token, perplexity 13.81): At equal KV cache footprint, the CLA2 model improves perplexity by 0.21 points. This is a 2× memory reduction relative to H128-MQA (10,240 bytes, 13.54) with only 0.06 points of perplexity degradation. The CLA2 model retains the larger head dimension (128 vs. 64) made possible by halving the number of KV caches.
H90-MQA-CLA2 (3,600 bytes/token, perplexity 13.73) vs. H46-MQA (3,680 bytes/token, perplexity 13.96): At comparable memory (3,600 vs. 3,680, with CLA2 actually using slightly less), the CLA2 model achieves 0.23 points better perplexity.
H64-MQA-CLA2 (2,560 bytes/token, perplexity 13.89) vs. H32-MQA (2,560 bytes/token, perplexity 14.37): At the smallest common memory footprint, the CLA2 model wins by 0.48 points — the largest absolute improvement among the pairwise comparisons.
H256-MQA-CLA2 (10,240 bytes/token, perplexity 13.51) vs. H128-MQA (10,240 bytes/token, perplexity 13.54): At the memory footprint of a standard d_head = 128 MQA model, CLA2 with larger heads achieves a marginal 0.03 point improvement. The paper notes this is a small gain, but it means CLA2 is never worse than the equal-memory baseline.
H512-MQA-CLA2 (20,480 bytes/token, perplexity 13.49) vs. H128-GQA2 (20,480 bytes/token, perplexity 13.52): At higher memory budgets where GQA2 becomes the relevant baseline, CLA2 achieves essentially identical perplexity (0.03 points better) while using the same memory as a GQA2 model with half the head dimension.

The Pareto frontier interpretation. The paper's central visual claim in Figure 3 is that "CLA enables accuracy/memory Pareto improvements relative to existing Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) architectures." The red (CLA) points lie to the left of and below the blue (non-CLA) points across most of the memory range, meaning that for a given perplexity target, CLA models require less KV cache memory, and for a given memory budget, CLA models achieve better perplexity. The frontier is most clearly advanced in the 2,500–5,000 bytes/token range, where the CLA2+MQA models with d_head ∈ {64, 90, 128} visibly separate from the plain MQA ladder.

GQA+CLA2 Ablation: CLA Benefits Are Specific to MQA (Table 1)

The paper trains three models combining GQA with CLA2 to test whether the benefits generalise beyond the MQA regime:

H256-GQA4-CLA2 (40,960 bytes/token, perplexity 13.38): Matched against the non-CLA baseline at the same memory, H128-GQA4 (40,960 bytes/token, 13.36), the CLA2 variant is essentially tied — a 0.02 point difference in favor of the baseline. There is no Pareto improvement.
H128-GQA4-CLA2 (20,480 bytes/token, perplexity 13.48): At this memory budget, the relevant non-CLA comparison is H128-GQA2 (20,480 bytes/token, 13.52). CLA2 achieves 0.04 points better perplexity, a marginal gain.
H128-GQA2-CLA2 (10,240 bytes/token, perplexity 13.59): Matched against H256-MQA-CLA2 at the same 10,240 bytes (13.51), the GQA2 variant is 0.08 points worse. Against H128-GQA2 (which has 20,480 bytes — 2× the memory), CLA2 saves 2× memory at a cost of 0.07 perplexity points, but the fair comparison at equal memory is to H256-MQA-CLA2, which wins.

The key finding: only the GQA2-CLA2 configuration "was able to achieve a perplexity better than the corresponding baseline model with the same KV cache footprint" (Section 3.2.1), and even then, MQA-CLA2 at the same memory outperforms it. The paper's conclusion is that CLA "appears to deliver the most robust benefits when used in conjunction with MQA" — the interaction between layer sharing and head sharing is not additive; combining both forms of sharing (GQA reduces heads, CLA reduces layers) yields diminishing returns compared to concentrating the memory reduction on the layer dimension (via CLA) while keeping per-layer head capacity maximized (via MQA).

Two models test whether more aggressive layer sharing (CLA3, CLA4) produces better accuracy/memory tradeoffs than CLA2:

H128-MQA-CLA3 (3,584 bytes/token, perplexity 13.77): At 3,584 bytes, the CLA3 model should be compared to the CLA2 models bracketing this memory range — H128-MQA-CLA2 at 5,120 bytes (13.60, less memory-constrained) and H90-MQA-CLA2 at 3,600 bytes (13.73, essentially equal memory). CLA3 is 0.04 points worse than H90-MQA-CLA2 at nearly identical memory (3,584 vs. 3,600), and substantially worse than the 5,120-byte CLA2 model. Against the non-CLA baseline H46-MQA (3,680 bytes, 13.96), CLA3 does show improvement (0.19 points), confirming it still Pareto-dominates plain MQA — it just does not dominate CLA2.
H128-MQA-CLA4 (2,560 bytes/token, perplexity 13.95): At 2,560 bytes, the comparison is to H64-MQA-CLA2 (2,560 bytes, 13.89) and the non-CLA H32-MQA (2,560 bytes, 14.37). CLA4 is 0.06 points worse than CLA2 at the same memory, but 0.42 points better than the plain MQA baseline. Again, CLA4 Pareto-dominates MQA but not CLA2.

The paper's conclusion (Section 4): "using sharing factors greater than 2 (CLA3 and above) achieves slightly worse accuracy/memory tradeoffs than using CLA2 and varying the head dimension." The optimal strategy for a given memory target is to set the sharing factor to 2 and adjust d_head to hit the memory budget, rather than increasing the sharing factor.

Three alternative CLA2 configurations test whether layers at the boundaries or concentrated blocks benefit from non-uniform sharing, all using MQA with d_head = 128:

H128-MQA-CLA2-KeepEnds (5,632 bytes/token, perplexity 13.62): Gives layers 0 and 19 their own dedicated KV caches, with uniform pairing in between. This uses more memory than uniform CLA2 (5,632 vs. 5,120 bytes, a 10% increase) and achieves 0.02 points worse perplexity (13.62 vs. 13.60). The marginal memory cost of keeping the ends is not recovered by improved accuracy.
H128-MQA-CLA2-DenseFront (5,632 bytes/token, perplexity 13.75): Concentrates KV-producing layers in the first 10 layers, with all subsequent layers reusing the KV cache of layer 9. This is 0.15 points worse than uniform CLA2 at higher memory cost — a clear regression.
H128-MQA-CLA2-DenseBack (5,632 bytes/token, perplexity 14.03): Concentrates KV-producing layers at the end, with early layers reusing the KV cache of layer 1. This is 0.43 points worse than uniform CLA2, the worst CLA variant tested, and worse than some non-CLA baselines at comparable memory (H46-MQA achieves 13.96 at 3,680 bytes — substantially less memory).

The non-uniform patterns all underperform uniform CLA2 despite requiring more memory (because L_KV = 11 instead of 10). The uniform pairing — where each KV cache is shared by exactly two consecutive layers with the producer being the earlier (even-indexed) layer — is the best configuration tested.

1B-Scale Learning Rate Tuning: CLA Benefits Survive Optimization (Tables 3, 4)

The learning rate tuning experiments on three key 1B-scale models verify that the design space exploration findings are not artifacts of a single suboptimal learning rate:

H128-MQA optimal LR: 1.5 × 10^{-3}, validation perplexity: 12.39
H128-MQA-CLA2 optimal LR: 2.25 × 10^{-3}, validation perplexity: 12.43
H64-MQA optimal LR: 2.25 × 10^{-3}, validation perplexity: 12.74

The CLA2 model achieves: (a) a 0.04 point perplexity degradation relative to the 2×-memory H128-MQA baseline (12.43 vs. 12.39) — smaller than the 0.06-point gap in the untuned comparison — and (b) a 0.31 point improvement over the equal-memory H64-MQA baseline (12.43 vs. 12.74) — larger than the 0.21-point gap in the untuned comparison. Both margins actually widen in CLA's favor after tuning.

On Wikitext (Table 3), the pattern is even stronger: H128-MQA-CLA2 achieves 19.29, which is 0.01 points better than H128-MQA (19.30) and 0.71 points better than H64-MQA (20.00). At this out-of-distribution evaluation, the CLA2 model effectively matches the 2×-memory baseline while halving its KV cache footprint.

Downstream benchmarks (Table 4) show no consistent winner: all three models are within 1–5 percentage points of each other on all seven tasks, with no systematic advantage for any architecture. For instance, H128-MQA scores highest on BoolQ (57.40 vs. 53.21 for CLA2 and 55.81 for H64-MQA), while CLA2 scores highest on OpenBookQA (21.4 vs. 19.0 and 19.4). This is consistent with the interpretation that CLA does not fundamentally change the model's task-level capabilities — it achieves comparable quality with less memory.

3B-Scale Experiments: CLA Benefits Persist, with Complications (Tables 5–8)

The 3B-scale experiments are reported in two sets, with the second set conducted on different hardware due to logistical constraints.

First set (d_head = 128 regime, Tables 5–6): Three models trained on 100B tokens with tuned learning rates:

H128-MQA optimal LR: 6.75 × 10^{-4}, validation perplexity: 9.52, Wikitext: 13.63
H128-MQA-CLA2 optimal LR: 2.25 × 10^{-3}, validation perplexity: 9.34, Wikitext: 13.25
H64-MQA optimal LR: 1.0 × 10^{-3}, validation perplexity: 9.48, Wikitext: 13.49

This result is unexpected: the CLA2 model (9.34) outperforms both baselines, including the 2×-memory H128-MQA (9.52). It achieves better perplexity with half the KV cache. Even more surprisingly, H64-MQA (9.48) outperforms H128-MQA (9.52) despite having half the KV cache — the opposite of the 1B-scale trend where smaller heads monotonically increased perplexity.

The authors do not fully explain this anomaly but acknowledge it: "we observed a result different than we had expected: at 3B scale, our MQA-CLA2 model achieves substantially better perplexities than both our d_head=128 and d_head=64 MQA baselines" (Section 3.3). A plausible interpretation is that the larger model has sufficient representational capacity that the head dimension bottleneck at d_head=128 is not binding, and the CLA model benefits from having fewer distinct weight matrices to optimize, concentrating the model's capacity. Alternatively, the d_head=128 MQA baseline may be undertrained or suboptimally tuned despite the learning rate sweep. The downstream benchmarks (Table 6) again show no consistent winner, with all models within a few percentage points of each other.

Second set (d_head = 64 regime, Tables 7–8): Because H64-MQA was the stronger baseline in the first set, the authors re-centered the comparison around d_head=64. This set was trained on a different cluster with a different data order, including a retrained H64-MQA-CLA2 to control for environment differences. All models used LR = 1.0 × 10^{-3} (additional learning rates tested were worse).

H64-MQA (8,192 bytes/token): Wikitext perplexity 12.94
H64-MQA-CLA2 (4,096 bytes/token): Wikitext perplexity 12.99
H32-MQA (4,096 bytes/token): Wikitext perplexity 13.34

This replicates the 1B-scale pattern exactly: CLA2 incurs 0.05 points of perplexity degradation relative to the 2×-memory H64-MQA baseline, and achieves 0.35 points improvement over the equal-memory H32-MQA baseline. The margin (0.35) is comparable to the 1B-scale Wikitext improvement (0.71, Table 3) though smaller, confirming that CLA2 delivers consistent gains at 3B scale when the baseline comparison is properly centered on head dimensions where the memory tradeoff is active.

Downstream benchmarks (Table 8) again show no consistent winner among the three models on the seven tasks.

Ablation Studies and Robustness Checks

GQA + CLA2 combination: Training H256-GQA4-CLA2, H128-GQA4-CLA2, and H128-GQA2-CLA2 reveals that the benefits of CLA are not additive with GQA's head-grouping — only GQA2-CLA2 improves over the equal-memory baseline (13.59 vs. the relevant MQA-CLA2 at 10,240 bytes, 13.51, and vs. the equal-memory H256-MQA-CLA2 which is better), while GQA4-CLA2 configurations are at best tied with their non-CLA GQA counterparts (Table 1). This establishes a boundary condition: CLA is most effective when per-layer KV head count is already minimized (MQA), and combining it with GQA yields diminishing returns.

Sharing factor >2: MQA-CLA3 (13.77 at 3,584 bytes) and MQA-CLA4 (13.95 at 2,560 bytes) both Pareto-dominate plain MQA but are worse than CLA2 matched at the same memory by adjusting head dimension (H90-MQA-CLA2 at 13.73, H64-MQA-CLA2 at 13.89). The finding that CLA2 dominates higher sharing factors is robust across the memory range tested (Table 1).

Non-uniform sharing patterns: CLA2-KeepEnds (13.62 at 5,632 bytes), CLA2-DenseFront (13.75 at 5,632 bytes), and CLA2-DenseBack (14.03 at 5,632 bytes) all underperform uniform CLA2 (13.60 at 5,120 bytes) despite using more memory. DenseBack is the worst CLA configuration tested, worse even than H46-MQA (13.96 at 3,680 bytes), confirming that concentrating KV production at the end of the model is particularly harmful (Table 1).

Learning rate robustness at 1B scale: Sweeping learning rates for H128-MQA, H64-MQA, and H128-MQA-CLA2 (Appendix A, Figure 5) shows that the CLA2 model's optimal LR (2.25 × 10^{-3}) is higher than that of the equal-head-dimension baseline H128-MQA (1.5 × 10^{-3}) but equal to the equal-memory baseline H64-MQA (2.25 × 10^{-3}). At their respective optimal LRs, the CLA2 model achieves 12.43 validation perplexity vs. 12.39 for H128-MQA (0.04 points worse, 2× memory savings) and 12.74 for H64-MQA (0.31 points better, equal memory). This confirms that the design space exploration findings are not artifacts of a learning rate that happened to favor CLA (Table 3).

Learning rate robustness at 3B scale: The first-set 3B tuning (Appendix A, Figure 4) finds optimal LRs of 6.75 × 10^{-4} for H128-MQA, 2.25 × 10^{-3} for H128-MQA-CLA2, and 1.0 × 10^{-3} for H64-MQA. The CLA2 model requires a 3.3× higher learning rate than the equal-head-dimension baseline — a large shift. At these optimal LRs, CLA2 (9.34) actually outperforms H128-MQA (9.52), a result stronger than at 1B scale. The second-set 3B experiments use a fixed LR of 1.0 × 10^{-3} for all models (other LRs tested were worse), finding the standard 1B-scale pattern (Tables 7–8).

Downstream benchmark consistency: Across all three sets of tuned-model evaluations (Tables 4, 6, 8), no architecture consistently wins or loses across tasks. Models are within a few percentage points on all benchmarks, suggesting that CLA does not systematically improve or degrade task-specific capabilities beyond what perplexity improvements would predict.

Training infrastructure robustness: The second-set 3B experiments were conducted "on a different training cluster using a different training software stack and data order," including a retrained H64-MQA baseline to control for environment differences. The finding that H64-MQA-CLA2 (12.99 Wikitext PPL) outperforms H32-MQA (13.34) by 0.35 points while matching H64-MQA (12.94) within 0.05 points confirms that CLA benefits are not specific to a particular training setup (Table 7).

Critical Assessment

Claim 1: "CLA provides a Pareto improvement over the memory/accuracy tradeoffs which are possible with traditional MQA."

This claim is well-supported for the specific regime tested: MQA-CLA2 with d_head ∈ {64, 90, 128} at 1B scale consistently achieves better perplexity than equal-memory MQA baselines with smaller heads (Figure 3, Table 1). The 3B-scale replication in the second set (Table 7) confirms the pattern. However, the claim's scope is narrower than the paper's abstract suggests in several respects:

It holds specifically for MQA combined with CLA2. The GQA+CLA2 results are weaker (GQA4-CLA2 is tied, GQA2-CLA2 is slightly worse than MQA-CLA2 at equal memory), and the paper does not test CLA with full MHA at all. The claim should be understood as "MQA-CLA2 Pareto-dominates plain MQA," not "CLA universally improves attention architectures."
It holds in the 1B–3B parameter, 20–32 layer, 2,048-token context regime. There is no evidence at larger scales (7B, 13B, 70B+), longer contexts (where KV cache pressure is most acute), or with different layer counts. The paper's choice of 20 and 32 layers is natural for the 1B and 3B scales, but CLA's effectiveness may depend on model depth — deeper models have more layers to share across, potentially changing the optimal sharing factor or the sharing pattern.
The learning rate shift (CLA models needing 1.5–3.3× higher learning rates) means that the Pareto improvement is only attained when hyperparameters are re-tuned for CLA. A practitioner who naively applies a baseline's optimal learning rate to a CLA variant will see smaller or zero gains. This is not a flaw — the paper demonstrates that re-tuning works — but it means CLA is not a "drop-in" change; it requires hyperparameter re-optimization.

Claim 2: "CLA reduces the size of the KV cache by another 2× while maintaining nearly the same accuracy as unmodified MQA."

This is supported at both scales. At 1B, H128-MQA-CLA2 (5,120 bytes) achieves 12.43 vs. 12.39 for H128-MQA (10,240 bytes) — the 2× storage reduction costs 0.04 perplexity points. At 3B (second set), H64-MQA-CLA2 (4,096 bytes) achieves 12.99 vs. 12.94 for H64-MQA (8,192 bytes) — a 0.05 point cost. Both are plausibly "nearly the same accuracy."

However, "nearly the same accuracy" is measured only through perplexity. The downstream benchmark results (Tables 4, 6, 8) show no consistent winner or loser, but they also do not establish that CLA models match baseline accuracy in a statistically rigorous sense — the samples are small (500–1,000 questions per benchmark typically), and no confidence intervals are reported. A CLA model that is 0.05 perplexity points worse could be meaningfully worse on specific tasks that are underrepresented in the benchmark suite.

More importantly, the claim does not address sequence length scaling. All training and evaluation uses 2,048-token sequences. At longer sequences (the regime where KV cache pressure is most severe), the quality degradation from CLA might be amplified — the representational gap between a consumer layer's query and a producer layer's keys could accumulate over more tokens, or the model's ability to attend over long distances using "stale" keys (from an earlier layer) might degrade. The paper acknowledges this gap implicitly by not testing long-context scenarios and by stating in Section 4 that "end-to-end inference efficiency evaluations of large, long-context models employing CLA" is future work.

Claim 3: "CLA should be used between pairs of consecutive layers, and CLA appears to deliver the most robust benefits when used in conjunction with MQA."

The first part (CLA2 is optimal) is well-supported within the tested range: CLA3 and CLA4 are worse than CLA2 at matched memory (Table 1). However, the paper only tests integer sharing factors (2, 3, 4). It does not test CLA2 with some layers unshared (hybrid CLA/non-CLA patterns beyond the three non-uniform variants), or dynamic sharing where the sharing factor varies by depth. The CLA2-KeepEnds result suggests that forcing all layers to share might be slightly suboptimal — the keep-ends variant uses marginally more memory but the perplexity difference (13.62 vs. 13.60) is within noise. A more systematic exploration of which specific layers benefit most from dedicated KV caches is not conducted.

The second part (CLA works best with MQA) is supported by the GQA+CLA2 results but is tested only at three GQA configurations (GQA4-CLA2 at two head dimensions, GQA2-CLA2 at one). The interaction space between n_group and L_KV is sparse. It's possible that an intermediate configuration (e.g., GQA2-CLA3) yields a better tradeoff at some memory target than either pure MQA-CLA2 or pure GQA, but this is not explored. The claim is empirically grounded but should be interpreted as "among the configurations we tested" rather than a proven global optimum.

Claim 4 (implicit): "CLA is orthogonal to MQA/GQA and can be combined with any of them."

This is architecturally true — the mechanisms are independent — but the empirical results show that orthogonality does not imply additive benefit. GQA4-CLA2 does not improve over GQA4 at equal memory; GQA2-CLA2 is worse than MQA-CLA2 at equal memory. So while CLA can be combined with GQA, the paper's own data suggests it should not be, at least at the scales and configurations tested. The abstraction of "orthogonality" is correct but practically misleading without the qualification that the combination yields diminishing returns.

Missing experiments that would strengthen the paper:

Long-context evaluation. All experiments use 2,048-token sequences. Testing at 8,192, 32,768, or longer contexts would address whether CLA's accuracy penalty grows with sequence length — a critical question for the use case the paper motivates (long-context serving).
Larger model scales. 1B and 3B parameters are small by contemporary standards (June 2024). Results at 7B or 13B would provide stronger evidence that CLA scales. The 3B anomaly (H128-MQA-CLA2 outperforming H128-MQA) raises questions about whether the 1B findings generalise monotonically — larger-scale experiments could resolve this.
Inference throughput benchmarks. The paper argues CLA's benefit is enabling larger batch sizes and longer sequences, but never measures actual inference throughput or latency. A simple experiment — measure maximum batch size before OOM with and without CLA, or measure total tokens/second at sequence length 8,192 — would ground the systems claims in concrete metrics.
Comparison to post-hoc compression combined with CLA. If CLA reduces the KV cache by 2×, can post-hoc quantization (KVQuant-style) compress it further? Do the techniques compose additively or do they interfere? This would help practitioners understand the full memory reduction stack.
Ablation on CLA's separate layernorm parameters. The paper mentions that CLA models use separate layernorm affine parameters for Q vs. KV projections (Section 3.1) but never ablates this choice. How much of CLA's performance depends on this degree of freedom versus simply sharing KV projections with shared layernorm?

The 3B anomaly deserves scrutiny. In the first 3B set, H64-MQA (9.48) outperforms H128-MQA (9.52) despite having half the KV cache. This is the opposite of the 1B-scale monotonic relationship where smaller heads = worse perplexity. It could indicate that the 100B-token training budget is insufficient for the larger-head model to converge (undertraining), that the learning rate sweep missed a better configuration for H128-MQA (the optimal LR search stopped when perplexity stopped improving, but the landscape might have local minima), or that at 3B scale with 32 layers, the head dimension of 128 is genuinely overparameterized for the amount of training data. The paper acknowledges the anomaly but does not resolve it. The second-set 3B experiments (re-centered on d_head=64) side-step the issue rather than explaining it. This means the paper's strongest 3B evidence (second set) is for the head dimension regime where the baseline comparison is cleanest, but the anomalous first-set result remains unexplained and could indicate that CLA's benefits interact with model scale in ways the paper does not fully characterize.

The "Pareto improvement" framing has a subtle limitation. The paper treats KV cache memory as the sole resource constraint and perplexity as the sole quality metric. In practice, model designers care about inference latency, training cost, parameter count, and downstream task performance. CLA2 reduces parameter count slightly and training FLOPs slightly (fewer weight matrices), but does not improve per-token generation latency. A Pareto improvement in the (memory, perplexity) plane does not guarantee a Pareto improvement in the multi-objective space that practitioners actually optimize over. The paper is transparent about this (Section 2.3 enumerates what CLA does and does not affect), but the headline "Pareto improvement" language in the abstract and conclusion should be understood as scoped to the memory/accuracy tradeoff specifically.

6. Limitations and Trade-offs

6.1 All Experiments Use Short (2,048-Token) Sequences; No Evidence CLA Works at Long Contexts

The assumption or constraint. Every model in the paper is trained and evaluated with a sequence length of 2,048 tokens (Section 3.1, Table 2). The paper's entire motivation for CLA — reducing KV cache memory to enable longer sequences and larger batch sizes — is built on long-context scenarios, yet no experiment exceeds 2,048 tokens. The authors explicitly defer this to future work in Section 4:

"We leave end-to-end inference efficiency evaluations of large, long-context models employing CLA as an interesting problem for future work."

The consequence. The KV cache pressure that CLA is designed to relieve is most acute at sequence lengths of 8k, 32k, 128k, or more — exactly the regime where the paper provides no data. There are at least two plausible failure modes that cannot be ruled out:

Accumulated representational mismatch. At 2,048 tokens, a KV-consumer layer uses keys and values from an earlier layer's hidden states. As sequence length grows, each attention operation covers a wider temporal span. The representational gap between a consumer's query (derived from layer ℓ's hidden state after ℓ blocks of processing) and a producer's keys (derived from layer p's hidden state after only p blocks) may compound over more distant token positions — the model might successfully attend to nearby tokens using "stale" keys but struggle to retrieve information from tokens far in the past, where the representational mismatch is more consequential.
Length-generalisation failure of the sharing pattern. The uniform CLA2 pattern was optimised empirically on 2,048-token sequences. It is possible that at longer sequences, non-uniform sharing (e.g., dedicated KV caches at layers responsible for long-range retrieval, or sharing patterns that vary with depth) becomes necessary, and the CLA2 sweet spot degrades.

What evidence exists in the paper. None. Every perplexity number, every downstream benchmark evaluation, and every learning rate sweep is at sequence length 2,048. The paper does not mention any pilot experiment at longer contexts or provide any theoretical argument for why CLA should scale to long sequences.

Mitigation status. Not addressed. The paper acknowledges the gap as future work (Section 4) but does not attempt even a small-scale long-context probe (e.g., evaluating one CLA2 model at 4,096 or 8,192 tokens on perplexity). This is the single most significant limitation for practitioners considering CLA for the use case the paper itself motivates — long-context inference.

6.2 The 3B-Scale Results Contain an Unexplained Anomaly That Undermines Confidence in Scaling Behaviour

The assumption or constraint. The paper presents CLA as scaling consistently from 1B to 3B parameters, but the first set of 3B-scale experiments (Tables 5–6) produces a result that contradicts the 1B-scale trend and is never resolved. At 1B scale, smaller head dimensions monotonically increase perplexity (H128-MQA: 13.54 → H64-MQA: 13.81 → H46-MQA: 13.96 → H32-MQA: 14.37). At 3B scale in the first set, this relationship inverts: H64-MQA (9.48 validation perplexity) outperforms H128-MQA (9.52) despite having half the KV cache. Even more strikingly, H128-MQA-CLA2 (9.34) outperforms both, achieving better perplexity with half the KV cache of H128-MQA.

The consequence. There are several possible explanations, none confirmed:

The H128-MQA baseline at 3B scale may be undertrained on 100B tokens — the larger head dimension creates more parameters that require more data to converge, and the 100B-token budget may be insufficient, artificially depressing the baseline's performance.
The learning rate sweep (which stopped when perplexity stopped improving) may have converged to a suboptimal local minimum for H128-MQA, despite the paper's protocol.
CLA may genuinely interact with model scale in a non-monotonic way — the benefits may be larger at 3B than at 1B for reasons the paper does not characterise.

The consequence for practitioners is uncertainty: if you train a 7B or 13B CLA model, do you expect the 1B-scale pattern (small degradation vs. 2×-memory baseline, clear win vs. equal-memory baseline), the 3B first-set pattern (CLA outperforms the 2×-memory baseline), or something else? The paper's decision to re-center the second 3B set on d_head = 64 (Tables 7–8) and report the standard pattern there means the paper's headline conclusion — CLA2 works — is supported, but the anomalous first set remains unexplained and could indicate hidden scaling dynamics that matter at larger scales.

What evidence exists in the paper. Tables 5 and 6 document the anomaly. The authors acknowledge it in Section 3.3:

"we observed a result different than we had expected: at 3B scale, our MQA-CLA2 model achieves substantially better perplexities than both our d_head=128 and d_head=64 MQA baselines. Moreover, our d_head=64 MQA baseline model achieves better perplexities than our tuned d_head=128 MQA baseline, despite having only 1/2 as much KV cache capacity."

The second-set experiments (Tables 7–8) re-establish the 1B-scale pattern but do not explain the first-set anomaly. The training infrastructure changed between sets (different cluster, different data order), which introduces a confound: the anomaly could be an artifact of the first environment, the second environment, or a genuine scaling phenomenon.

Mitigation status. Partially addressed. The second-set experiments show that CLA2 works as expected when centered on d_head = 64, which is the head dimension regime most relevant to memory-constrained deployment. But the paper does not resolve whether the H128-MQA underperformance at 3B is a training budget issue, a hyperparameter issue, or an architecture-scaling issue. A simple additional experiment — training H128-MQA at 3B for 200B tokens to check if the perplexity improves relative to H64-MQA — would have clarified whether undertraining is the cause, but it is not performed.

6.3 CLA Models Require Re-Tuned Learning Rates; Naïve Transfer of Baseline Hyperparameters Underestimates CLA Performance

The assumption or constraint. The paper demonstrates that CLA models consistently require higher learning rates than their non-CLA counterparts with the same head dimension. At 1B scale, H128-MQA-CLA2's optimal LR is 2.25 × 10^{-3} versus 1.5 × 10^{-3} for H128-MQA — a 1.5× increase (Table 3). At 3B scale, the gap is 3.3×: 2.25 × 10^{-3} for H128-MQA-CLA2 versus 6.75 × 10^{-4} for H128-MQA (Table 5). The paper's own results would be substantially weaker if CLA models were evaluated at the baseline's learning rate: Figures 4 and 5 in Appendix A show that at LR = 1.5 × 10^{-3}, H128-MQA-CLA2 is still near its optimum at 1B scale, but at 3B scale, using H128-MQA's optimal LR of 6.75 × 10^{-4} for the CLA2 model (instead of its actual optimum 2.25 × 10^{-3}) would substantially degrade CLA2's reported perplexity.

The consequence. CLA is not a drop-in architectural change. A practitioner replacing MQA layers with CLA2 must re-execute the learning rate sweep, which costs additional compute. If the practitioner cannot afford a full sweep and simply reuses the baseline's learning rate, the CLA model's accuracy will be worse than the paper's headline numbers — potentially erasing the Pareto improvement entirely at larger scales where the learning rate gap is largest. The paper does not quantify how much CLA performance degrades at the "wrong" learning rate (the learning rate sweep curves in Appendix A provide this implicitly but are not discussed in these terms).

More subtly, this finding implies that other hyperparameters may also need re-tuning: weight decay, warmup duration, batch size, and potentially even architectural choices like layer count or FFN ratio. The paper only tunes the learning rate, so the full hyperparameter sensitivity of CLA models is unknown. If CLA shifts the optimal weight decay or warmup as it shifts the optimal learning rate, then the reported perplexities — even after LR tuning — may not represent the best achievable CLA performance.

What evidence exists in the paper. Tables 3, 5, and 7 report the optimal learning rates, showing consistent upward shifts for CLA models. Appendix A provides the full learning rate sweep curves visually (Figures 4 and 5), confirming that the optimal LR differs. The paper notes the pattern in Section 3.2.2 ("we found preliminary evidence to suggest that CLA models benefit from training with higher learning rates than comparable non-CLA models") but treats it as an observation rather than a limitation.

Mitigation status. The paper addresses the learning rate tuning threat by actually performing the sweeps for key comparisons (Sections 3.2.2, 3.3), which is methodologically rigorous. However, it does not provide guidance on how to set the learning rate for a new CLA configuration without a full sweep, does not explore whether the shift is predictable (e.g., as a function of the fraction of parameters removed), and does not tune other hyperparameters. A practitioner reading the paper would know they need to re-tune the LR but would not know whether weight decay, warmup, or other settings also need adjustment.

6.4 Single Model Family, Single Dataset, Single Task Modality; No Evidence of Generalisation Beyond SlimPajama Language Modelling

The assumption or constraint. All experiments train decoder-only Llama-like transformers on the SlimPajama dataset and evaluate primarily on language modelling perplexity (validation SlimPajama, Wikitext). The architecture, data, and task are held constant: there are no experiments on other model families (e.g., encoder-decoder, mixture-of-experts), other datasets (e.g., code, scientific text, multilingual corpora), other modalities, or other tasks (e.g., fine-tuning on downstream applications, instruction following, generation quality as judged by humans or LLM evaluators).

The consequence. The paper's claims — that CLA2+MQA advances the accuracy/memory Pareto frontier, that CLA2 should be the default sharing factor, that CLA works best with MQA — may be specific to the SlimPajama + Llama-like architecture combination. Several failure modes are plausible:

Domain shift. SlimPajama is predominantly web text. Code generation, mathematical reasoning, or multilingual text may have different redundancy patterns across layers. It is possible that in highly structured domains (code), adjacent layers learn more distinct representations, making cross-layer sharing more costly.
Architecture interaction. Rotary position embeddings (RoPE) encode position by rotating query and key vectors. When a KV-consumer layer uses keys from an earlier layer, those keys were position-encoded at the earlier layer's hidden state. The interaction between cross-layer KV sharing and RoPE — specifically, whether the position information embedded in the producer's keys remains interpretable to the consumer's queries — is not analysed. If RoPE's effectiveness depends on keys and queries being derived from the same hidden state, CLA could degrade position-aware attention.
Fine-tuning vs. pretraining. The paper only evaluates pretraining perplexity and zero-shot downstream benchmarks. It is unknown whether CLA models fine-tune as effectively as non-CLA models on downstream tasks — the shared KV representations might limit the model's ability to adapt to task-specific attention patterns during fine-tuning.

What evidence exists in the paper. The downstream benchmark evaluations (Tables 4, 6, 8) provide some evidence beyond perplexity, but they are zero-shot evaluations of pretrained models, not fine-tuned task performance. The benchmarks cover commonsense reasoning and question answering, which are narrow compared to the range of tasks LLMs are deployed for (code generation, dialogue, summarisation, translation, retrieval-augmented generation). No experiment varies the pretraining data distribution, the model architecture family, or the position encoding scheme.

Mitigation status. Not addressed. The paper's scope is explicitly limited to pretraining experiments on SlimPajama with a Llama-like architecture, and it does not claim generalisation beyond this setting. However, the abstract states CLA "enables inference with longer sequence lengths and larger batch sizes than would otherwise be possible" without qualifying that this has only been shown for SlimPajama language modelling with Llama-like architectures at short context lengths. A practitioner using a different model family (e.g., a Mixture-of-Experts architecture, or a non-RoPE position encoding) or a different domain (code, multilingual text) cannot assume the CLA2 recipe transfers without degradation.

6.5 CLA Does Not Reduce Per-Token Generation Latency; Memory Savings Enable Throughput Gains Only If Memory Was the Bottleneck

The assumption or constraint. As the paper states in Section 2.3, CLA provides no direct latency benefit to the core attention computation:

"Unlike MQA and GQA, CLA has no direct effect on the memory bandwidth consumed by the attention mechanism in each decoding step, because even shared KV cache layers must be separately re-read from main memory in each attention layer. CLA therefore has no direct effect on the latency of the core attention computation during decoding."

The latency benefits of CLA are entirely indirect: by reducing KV cache memory, CLA enables larger batch sizes (which improves throughput but not per-request latency) or enables serving requests with longer sequences that would otherwise exceed memory capacity.

The consequence. CLA is a memory capacity expansion technique, not a latency optimisation. For deployment scenarios where memory is abundant but per-token generation speed is the bottleneck — e.g., serving a single long-sequence request on a GPU with excess HBM — CLA provides no benefit. The memory savings from CLA translate to throughput improvements only if the deployment was previously memory-bound: that is, the maximum batch size was limited by KV cache memory rather than by compute or memory bandwidth. In compute-bound or bandwidth-bound regimes, reducing KV cache memory does not increase the number of requests that can be processed simultaneously, because the GPU is already saturated with a smaller batch.

The paper's systems discussion (Section 2.3) is transparent about this, but the abstract and introduction frame CLA primarily in terms of memory reduction without clarifying the limited translation to latency. A practitioner reading only the abstract ("enabling inference with longer sequence lengths and larger batch sizes") might assume CLA provides general inference acceleration, which it does not.

What evidence exists in the paper. Section 2.3 explicitly distinguishes between KV cache memory, training memory, parameters/FLOPs, and core attention latency. However, no experiment measures actual inference throughput or latency — there are no tokens-per-second measurements, no maximum-batch-size-before-OOM benchmarks, and no comparison of CLA vs. non-CLA serving performance at long sequences. The paper's systems claims are architectural arguments, not empirical measurements.

Mitigation status. The paper is transparent about the scope of CLA's latency effects in Section 2.3 but does not empirically validate even the indirect throughput claims. The future work section (Section 4) acknowledges this: "We leave end-to-end inference efficiency evaluations... as an interesting problem for future work." A practitioner deploying CLA would need to run their own inference benchmarks to determine whether the memory savings actually translate to throughput improvements in their specific hardware and workload configuration, and the paper provides no data to guide that expectation.

6.6 Only 1B and 3B Parameter Scales Tested; No Evidence at the 7B–70B+ Scales Where KV Cache Pressure Is Most Acute in Production

The assumption or constraint. The paper's largest model has 3 billion parameters and 32 layers, trained on 100 billion tokens. Contemporary production LLMs (as of early 2024) routinely operate at 7B, 13B, 34B, 70B, or larger parameter counts, often with 40–80+ layers and trained on trillions of tokens. These models are precisely where KV cache memory pressure is most severe — a 70B model at 128k context can require tens of gigabytes of KV cache memory, dwarfing the model weights in some configurations. Yet the paper provides no evidence that CLA's accuracy/memory tradeoffs hold at these scales.

The consequence. Scaling laws for architectural modifications are not guaranteed to be monotonic. Several patterns in the paper's own data suggest non-trivial scale dependence:

The optimal learning rate shift between CLA and non-CLA models is 1.5× at 1B but 3.3× at 3B (Tables 3 and 5) — the gap grows with scale, suggesting the optimisation dynamics change non-linearly.
The first-set 3B anomaly (H64-MQA outperforming H128-MQA, H128-MQA-CLA2 outperforming both) indicates that the relationship between head dimension, KV cache size, and perplexity shifts between 1B and 3B parameters.
CLA2's perplexity degradation relative to the equal-head-dimension baseline changes from 0.04 points at 1B (Table 3) to +0.18 points improvement at 3B (Table 5, first set) — the effect size is not stable.

At 70B parameters with 80 layers, the sharing dynamics may be fundamentally different. With more layers, the representational gap between adjacent layers may be smaller (because each layer performs a smaller transformation in a deeper model), potentially making CLA more effective — or the cumulative gap across 40 CLA2 pairs may be larger, making CLA less effective. The paper provides no basis for predicting which.

What evidence exists in the paper. The paper explicitly studies 1B and 3B scales, with the architectural and training hyperparameters in Table 2. The authors do not claim that their results extrapolate to larger scales. However, the practical target audience for KV cache reduction techniques is precisely the operators of large-scale LLMs, and the paper offers no evidence that the CLA2 recipe transfers.

Mitigation status. Not addressed beyond the scope limitation. The paper trains 1B and 3B models "from scratch," and scaling to 7B+ would require substantially more compute than is likely available in an academic setting — this is an understandable practical constraint. But the absence of larger-scale experiments means that the paper's primary practical recommendation ("we recommend this recipe to practitioners as a conservative change to existing MQA architectures," Section 4) is based on evidence at scales that are an order of magnitude smaller than typical production deployments. A conservative practitioner would want at least a 7B-scale validation before adopting CLA in a production 70B model.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper is best understood not as a paradigm shift but as a precise, empirically grounded reframing of the KV cache design space. Before CLA, the field implicitly treated the number of layers L in the KV cache size formula S_{KV} = L × 2 · n_group · d_head · p as a fixed architectural constant—every layer got its own KV cache because that was the default. MQA and GQA had established that the n_group term could be productively reduced, opening one dimension of the design space. CLA demonstrates that the L term is also a tunable knob, not a requirement, and that reducing it (specifically from L to L/2 via uniform CLA2) interacts favorably with existing head-sharing techniques—particularly MQA—to push the accuracy/memory Pareto frontier outward.

The conceptual shift is subtle but important. Prior work treated the KV cache memory problem as a question of compression: given a fixed architecture with per-layer KV caches, how can we reduce storage via quantization, eviction, or head-sharing? CLA instead treats the KV cache as an architectural budget that can be allocated across layers, not a mandatory per-layer expense. The insight that adjacent layers contain redundant KV representations—and that this redundancy can be designed out from scratch at training time rather than compressed post-hoc—reframes KV cache efficiency from a downstream optimization problem to an upstream architectural design choice.

The paper reconciles no major prior contradictions (there was no active debate about whether cross-layer sharing would work—the question had simply not been asked systematically), but it resolves an uncertainty that was latent in the field: could the success of MQA/GQA along the head dimension be replicated along the depth dimension? The answer is a qualified yes—qualified because it works robustly only at CLA2, only with MQA, and has been demonstrated only at 1B–3B scales with short contexts. The boundary conditions are as important as the affirmative result.

The work makes certain research directions more attractive and others less so:

More attractive: principled architecture co-design. The finding that CLA2 works best with MQA (not GQA) suggests that redundancy across heads and redundancy across layers are partially substitutable. This opens a research program to systematically characterize the joint (n_group, L_KV) design space, potentially discovering optimal allocation rules as a function of model scale, depth, and training budget. The paper's memory-matched comparison methodology provides the template for such exploration.
More attractive: training-time architectural interventions over post-hoc compression. CLA demonstrates that models trained from scratch with a reduced KV cache structure can outperform models that are first trained with full KV caches and then compressed. This does not invalidate post-hoc methods (which remain necessary for already-trained models), but it suggests that the next generation of models should incorporate KV cache constraints into the pretraining architecture rather than treating them as a serving-time afterthought. The CLA2+MQA recipe is a concrete, low-risk starting point for model designers.
Less attractive: aggressive cross-layer sharing beyond factor 2. The CLA3 and CLA4 results (Table 1) show diminishing or negative returns relative to CLA2 at matched memory. The representational gap between a KV consumer and a producer more than one layer away appears to be large enough that the model cannot fully compensate, even with larger head dimensions. This suggests that pushing layer sharing to extreme factors without additional mechanisms (e.g., learned transformations between producer and consumer KV representations) is not a productive direction.
Less attractive: treating MHA as the default attention architecture. CLA's strongest results are with MQA, and the paper shows GQA+CLA underperforms MQA+CLA at equal memory. Combined with the widespread adoption of GQA and MQA in production models (Llama, Mistral, Gemma), this strengthens the case that Multi-Head Attention—with its n_group = n_query default—represents an overparameterized baseline that is rarely optimal under memory constraints. The burden of proof is shifting toward justifying why a model should use MHA rather than MQA with CLA2.

Follow-Up Research This Work Enables

Long-context evaluation of CLA2+MQA at 8k–128k tokens with controlled perplexity and retrieval benchmarks. The most urgent gap in the paper is the absence of any experiment beyond 2,048-token sequences. A strong follow-up would train (or continue-training from a CLA2 checkpoint) an H128-MQA-CLA2 and an H64-MQA baseline at 1B scale on sequences of 8,192 or 32,768 tokens, measuring validation perplexity as a function of position within the sequence. The key question is whether the representational gap between a consumer layer's query and a producer layer's keys accumulates over long contexts, causing perplexity degradation on tokens that require attending to distant positions. A complementary needle-in-a-haystack retrieval evaluation (e.g., placing a fact at position k and querying at the end of a 32k-token sequence) would probe whether CLA impairs long-range information retrieval specifically. If CLA2 shows no degradation at long contexts, the paper's motivating use case is validated; if it degrades, the result would bound CLA's applicability and motivate layer-specific sharing patterns optimized for retrieval depth.

Joint optimization of (n_group, L_KV) across model scales to establish CLA scaling laws. The paper explores a sparse grid of (n_group, L_KV) configurations at two model scales, but does not systematically map how the optimal configuration varies with total KV cache budget, model depth, and training tokens. A scaling-law study—analogous to Chinchilla for pretraining or the compute-optimal test-time scaling work—would train models at multiple scales (e.g., 300M, 1B, 3B, 7B) with multiple (n_group, L_KV) settings, each at multiple training token budgets, and fit a parametric model predicting perplexity as a function of n_group, L_KV, d_head, model size, and training FLOPs. The paper's finding that GQA4+CLA2 shows no benefit over GQA4 while MQA+CLA2 shows clear gains suggests the interaction is non-trivial and worth characterizing precisely. Such a study would also resolve the 3B anomaly (H64-MQA outperforming H128-MQA in the first set) by distinguishing undertraining effects from genuine architectural scaling trends.

CLA with learned KV transformations between producer and consumer layers. The paper's CLA2 design passes the KV-producer's keys and values unchanged to the consumer layer. This is the simplest form of sharing, but it forces the consumer's queries to adapt entirely to the producer's representational space. A natural extension is to insert a lightweight learned transformation between the producer's KV cache and the consumer's attention computation—for example, a small linear projection or a low-rank adapter applied to K_p and V_p before they are used by layer ℓ. This would give each consumer layer the ability to "translate" the producer's KV representations into its own representational space, potentially reducing the accuracy penalty of CLA and enabling larger sharing factors. The cost is a small increase in parameters and compute (one small projection matrix per consumer layer), which would need to be weighed against the memory savings. A concrete experiment: compare H128-MQA-CLA2 with and without learnable 32-dimensional bottleneck adapters on the shared KVs at 1B scale, measuring whether the perplexity gap to the 2×-memory baseline shrinks.

Stress-testing CLA against architecture variations: non-RoPE position encodings, MoE, encoder-decoder. The paper uses a Llama-like architecture with RoPE throughout. RoPE encodes position information directly into keys and queries via rotation. When a CLA consumer layer attends over keys from an earlier layer, those keys carry position information encoded at the producer's hidden state. It is unknown whether RoPE's effectiveness depends on keys and queries being derived from the same hidden state, and whether alternative position encoding schemes (ALiBi, learned absolute positions, no positional encoding) interact differently with CLA. A controlled experiment training 1B models with H128-MQA-CLA2 under RoPE, ALiBi, and learned positional embeddings would reveal whether CLA's effectiveness is tied to the position encoding choice. Similarly, applying CLA in a Mixture-of-Experts architecture (where different tokens are routed through different FFN experts, potentially making layer representations more heterogeneous) or in an encoder-decoder setting (where the encoder's KV cache is separate from the decoder's) would test the generality of the CLA2 sweet spot.

Combining CLA with post-hoc KV cache quantization to measure composability. The paper positions CLA as orthogonal to post-hoc compression, but never tests the combination. A practical deployment would likely use both: CLA2 to halve the architectural KV cache size, then KVQuant-style quantization to reduce precision to 2–4 bits per element. The key question is whether CLA's shared KV representations are more or less amenable to quantization than per-layer unique caches. If CLA produces KV representations that are used by multiple layers (each with potentially different sensitivity to quantization error), the optimal quantization scheme might need to be layer-aware. A concrete experiment: apply KVQuant-style per-channel quantization to the KV caches of H128-MQA and H128-MQA-CLA2 at 1B scale, measuring perplexity degradation at 2-bit and 4-bit precision for each model. If CLA+KVCquant at 2-bit achieves similar perplexity to plain MQA at 4-bit, the total effective memory reduction is 4× (2× from CLA, 2× from quantization) while maintaining accuracy—this would dramatically strengthen the practical case for architectural KV reduction.

Fine-tuning and continual pretraining of CLA models for downstream tasks. The paper evaluates CLA only on pretraining perplexity and zero-shot benchmarks. It is unknown whether CLA models—with their shared KV representations—fine-tune as effectively as non-CLA models on downstream tasks requiring task-specific attention patterns (e.g., instruction following, long-form summarization, multi-turn dialogue). The concern is that shared KV caches might limit the model's ability to develop layer-specialized attention patterns during fine-tuning. A follow-up would fine-tune the learning-rate-tuned 1B models (H128-MQA, H128-MQA-CLA2, H64-MQA) on a suite of tasks (e.g., Alpaca instruction-tuning, CNN/DailyMail summarization, a multi-turn dialogue dataset) and measure final task performance relative to the pretraining perplexity ranking. If CLA models fine-tune as well as their perplexity predicts, the pretraining results directly translate to downstream utility. If they underperform relative to perplexity, it would suggest a hidden cost of KV sharing that only emerges during adaptation.

Practical Applications and Downstream Use Cases

On-device and edge deployment of small LLMs with constrained memory. The 1B-scale results are directly applicable to scenarios where LLMs run on consumer devices (laptops, phones) with limited unified memory. An H128-MQA-CLA2 model requires 5,120 bytes/token for the KV cache versus 10,240 bytes/token for H128-MQA—a 2× savings. At a 4,096-token context with a batch size of 1, this reduces KV cache memory from roughly 42 MB to 21 MB at 16-bit precision. On a device with 8 GB of available memory also holding model weights (~2 GB for a 1B-parameter model in 16-bit), every megabyte saved on the KV cache can be allocated to longer context or to running other processes concurrently. The paper's recommendation—"use MQA-CLA2 with the largest head dimension that fits your memory budget"—provides a concrete deployment recipe: profile the target device's available memory, calculate the maximum bytes/token budget after accounting for weights and activations, and select d_head accordingly under the CLA2 formula L_KV × 2 × 1 × d_head × 2.

Batch inference pipelines where KV cache memory limits throughput. In server-side batch inference serving many requests simultaneously, KV cache memory is often the binding constraint on batch size when sequence lengths are long. The paper's 3B-scale second-set results (Table 7) show that H64-MQA-CLA2 (4,096 bytes/token) achieves Wikitext perplexity 12.99 versus H64-MQA at 8,192 bytes/token (12.94)—a 0.05-point degradation for a 2× memory reduction. For a deployment serving 3B-parameter models at 8,192-token sequences, halving the per-request KV cache means the same GPU memory can serve approximately twice as many requests in a batch, roughly doubling throughput, at the cost of a perplexity delta that is likely imperceptible in most applications. The key caveat from Section 2.3 applies: this throughput gain materializes only if the deployment was memory-bound (batch size limited by KV cache memory, not compute). Practitioners should profile their serving stack to determine whether KV cache memory is the active bottleneck before adopting CLA for throughput; if the bottleneck is compute or attention bandwidth, CLA's memory savings will not translate to throughput improvements.

Persistent KV caches for multi-turn conversations and long-prefix reuse. Systems like Google's Gemini context caching and AttentionStore (Gao et al., 2024) cache KV activations to avoid recomputing attention over shared prefixes across multiple requests or conversation turns. The storage cost of these persistent caches scales directly with the per-token KV cache size. CLA2's 2× reduction in per-token KV cache memory directly halves the storage footprint of persistent caches. For a service maintaining KV caches for millions of conversations, each with thousands of tokens of shared prefix, CLA2 halves the storage infrastructure cost. This is a pure storage win—the paper's finding that CLA provides no direct latency benefit (Section 2.3) is irrelevant for this use case because the bottleneck is storage capacity and retrieval bandwidth, not per-step attention latency. The 1B-scale Wikitext result (Table 3) is particularly encouraging: H128-MQA-CLA2 achieves 19.29 versus 19.30 for the 2×-memory H128-MQA, meaning the stored caches lose essentially no information despite being half the size.

Training-efficient architectures for organizations with limited compute budgets. The paper notes that CLA slightly reduces parameter count and training FLOPs (fewer W^K and W^V weight matrices). While this is a secondary benefit, it matters for research labs and startups that train models from scratch with constrained compute. At 1B scale with MQA-CLA2, removing W^K and W^V from 10 of the 20 layers eliminates 10 × 2 × d_model × d_head parameters—for d_model = 2048 and d_head = 128, roughly 5.2 million parameters, or about 0.5% of a 1B-parameter model. This is small but non-zero, and the training FLOPs savings are proportional. More importantly, the paper's finding that CLA2+MQA at equal memory outperforms plain MQA (e.g., H128-MQA-CLA2 vs. H64-MQA, a 0.21–0.31 perplexity improvement) means that organizations targeting a specific KV cache memory budget can achieve better pretraining quality within that budget by using CLA rather than shrinking head dimensions. This is effectively a free lunch: for a given deployment memory constraint, CLA delivers better models than the previously best-known architecture (plain MQA with small heads). The cost is the need to re-tune the learning rate, which the paper's Appendix A data suggests is manageable (sweeping LRs from the baseline's optimum upward in 1.5× increments until validation perplexity stops improving).

Dependency tree

Papers this one builds on, walking backwards through the same domain and shared topics.

GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints2023· architecture← LLM-judged similar
- Read More2023· architecture← in-domain predecessor
Layer-Condensed KV Cache for Efficient Inference of Large Language Models2024· architecture← in-domain predecessor
- KAN: Kolmogorov-Arnold Networks2024· architecture← in-domain predecessor
- ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition2024· serving← topic thread
Forgetting Transformer: Softmax Attention with a Forget Gate2025· architecture← similar topic
- Jamba: A Hybrid Transformer-Mamba Language Model2024· architecture← LLM-judged similar
- MoBA: Mixture of Block Attention for Long-Context LLMs2025· architecture← in-domain predecessor
Selective Attention Improves Transformer2024· architecture← similar topic
- Configurable Foundation Models: Building LLMs from a Modular Perspective2024· architecture← in-domain predecessor

What it influenced

Later papers in the same chain. Direction reversed.

Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality2024· architecture→ in-domain predecessor
- Block Transformer: Global-to-Local Language Modeling for Fast Inference2024· architecture→ in-domain predecessor
- ETC: Encoding Long and Structured Inputs in Transformers2020· architecture→ similar topic
- Hungry Hungry Hippos: Towards Language Modeling with State Space Models2022· architecture→ similar topic
- MambaByte: Token-free Selective State Space Model2024· architecture→ similar topic
Rethinking Attention: Exploring Shallow Feed-Forward Neural Networks as an Alternative to Attention Layers in Transformers2023· architecture→ similar topic
- Memory Augmented Language Models through Mixture of Word Experts2023· architecture→ in-domain predecessor
- Deep Sparse Rectifier Neural Networks2000· architecture→ similar topic
- EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks2019· architecture→ similar topic
- Learning to Execute Programs with Instruction Pointer Attention Graph Neural Networks2020· architecture→ similar topic
Selective Attention Improves Transformer2024· architecture→ similar topic
- Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention2025· architecture→ LLM-judged similar
- DIFFERENTIAL TRANSFORMER2024· architecture→ in-domain predecessor
- FLEX ATTENTION: A PROGRAMMING MODEL FOR GENERATING OPTIMIZED ATTENTION KERNELS2024· architecture→ similar topic