# EG-MLA: Embedding-Gated Multi-head Latent Attention for Scalable and Efficient LLMs
**Authors**: Zhengge CaiHaowen Hou
## Abstract
Reducing the key-value (KV) cache size is a crucial step toward enabling efficient inference in large language models (LLMs), especially under latency and memory constraints. While Multi-Head Attention (MHA) offers strong representational power, it incurs significant memory overhead. Recent work on Multi-head Latent Attention (MLA) mitigates this by compressing KV representations into a shared latent space, achieving a better trade-off between performance and cache efficiency. While MLA already achieves significant KV cache reduction, the scope for further compression remains limited without performance loss. In this paper, we propose Embedding-Gated Multi-head Latent Attention (EG-MLA), a novel extension of MLA that further reduces KV cache size while enhancing representational expressiveness. EG-MLA introduces a token-specific embedding gating mechanism applied in the latent space, enabling fine-grained modulation of compressed KV vectors with minimal additional computation. Compared to MHA, EG-MLA achieves over 91.6% reduction in KV cache size with negligible performance degradation. Relative to MLA, EG-MLA consistently improves task accuracy across diverse reasoning benchmarks while achieving up to 59.9% additional memory savings. Our theoretical analysis highlights how embedding gating induces implicit high-order interactions, and empirical evaluations demonstrate robust generalization across model scales and compression regimes. Notably, we successfully scale EG-MLA to over 1 billion parameters, demonstrating its practical viability for large-scale LLM deployment. These results establish EG-MLA as a memory- and compute-efficient attention mechanism that enables scalable, high-performance inference in modern LLMs.
## Introduction
<details>
<summary>x1.png Details</summary>
### Visual Description
## Line Chart: Average Accuracy vs Training Tokens (Quadrant a)
### Overview
A line chart comparing the average accuracy of multiple models across different training token scales (10B to 50B). The chart shows six distinct data series with varying performance trajectories.
### Components/Axes
- **X-axis**: Training Tokens (10B, 20B, 30B, 40B, 50B)
- **Y-axis**: Avg. Accuracy (41% to 45%)
- **Legend**:
- Green: MLA-256
- Blue: EG-MLA-emb256
- Orange: EG-MLA-emb128
- Purple: EG-MLA-emb512
- Red: EG-MLA-emb1024
- Pink: EG-MLA-emb5048
### Detailed Analysis
1. **MLA-256** (Green): Starts at ~42.2% at 10B, rises steadily to ~43.8% at 50B
2. **EG-MLA-emb256** (Blue): Begins at ~41.8% at 10B, plateaus at ~43.2% by 30B, then climbs to ~44.1% at 50B
3. **EG-MLA-emb128** (Orange): Starts at ~42.0% at 10B, follows a similar trajectory to EG-MLA-emb256 but with slightly lower gains
4. **EG-MLA-emb512** (Purple): Begins at ~42.5% at 10B, shows the steepest ascent to ~44.8% at 50B
5. **EG-MLA-emb1024** (Red): Starts at ~42.8% at 10B, maintains the highest accuracy throughout, reaching ~45.2% at 50B
6. **EG-MLA-emb5048** (Pink): Starts at ~42.3% at 10B, shows the most gradual improvement to ~45.5% at 50B
### Key Observations
- All models show improved accuracy with increased training tokens
- Larger embedding sizes (5048) achieve highest final accuracy but with slower improvement rate
- MLA-256 (smallest embedding) shows most consistent gains
- EG-MLA-emb1024 maintains lead throughout training progression
### Interpretation
The chart demonstrates that larger embedding sizes generally yield better performance, but with diminishing returns. The MLA-256 model shows optimal efficiency, achieving ~43.8% accuracy at 50B tokens with minimal computational overhead. The EG-MLA-emb1024 model, while achieving highest accuracy, requires significantly more resources for marginal gains.
---
## Line Chart: Loss vs Embedded Size (Quadrant b)
### Overview
A line chart comparing loss metrics across different embedded sizes (38M to 1236M) for two model configurations.
### Components/Axes
- **X-axis**: Embedded Size (38M, 77M, 154M, 309M, 618M, 1236M)
- **Y-axis**: Loss (3.05 to 3.30)
- **Data Series**:
- Red dashed line: MLA-Base
- Blue line: EG-MLA-Base
### Detailed Analysis
1. **MLA-Base** (Red dashed): Maintains constant loss of ~3.30 across all embedded sizes
2. **EG-MLA-Base** (Blue):
- Starts at ~3.32 at 38M
- Decreases to ~3.22 at 77M
- Drops to ~3.12 at 154M
- Sharp decline to ~3.05 at 309M
- Slight increase to ~3.07 at 618M
- Final drop to ~3.05 at 1236M
### Key Observations
- EG-MLA-Base shows significant loss reduction with larger embeddings
- Loss reduction plateaus at ~3.05 beyond 309M embedded size
- MLA-Base maintains stable but higher loss (~3.30)
### Interpretation
The EG-MLA-Base configuration demonstrates superior scalability, achieving 17% loss reduction compared to MLA-Base at maximum embedded size. The performance plateau at 309M suggests diminishing returns beyond this point, indicating optimal resource allocation at this scale.
---
## Bar Chart: KV Cache Efficiency (Quadrant c)
### Overview
A comparative bar chart showing KV cache efficiency (K/Token #Elem) for three attention mechanisms.
### Components/Axes
- **X-axis**: Attention Mechanisms (MHA, MLA-Base, EG-MLA-Base)
- **Y-axis**: KV Cache Elements (0-15)
- **Annotations**:
- EG-MLA-Base shows 59.9% reduction vs MLA-Base
- EG-MLA-Base shows 91.6% reduction vs MHA
### Detailed Analysis
1. **MHA**: Highest KV cache usage at ~15 elements
2. **MLA-Base**: ~5 elements (66.7% reduction vs MHA)
3. **EG-MLA-Base**: ~0.75 elements (95% reduction vs MHA)
### Key Observations
- EG-MLA-Base achieves dramatic KV cache efficiency improvements
- MHA requires 20x more elements than EG-MLA-Base
- MLA-Base shows intermediate efficiency gains
### Interpretation
The EG-MLA-Base configuration reduces KV cache requirements by nearly 92% compared to traditional MHA, demonstrating significant memory efficiency gains. This suggests potential for 10-15x memory savings in large-scale deployments.
---
## Bar Chart: Decoding Latency (Quadrant d)
### Overview
A comparative bar chart showing decoding latency (ms/Token) for two model configurations.
### Components/Axes
- **X-axis**: Model Configurations (MLA-Base, EG-MLA-Base)
- **Y-axis**: Decoding Latency (0-25 ms/Token)
- **Data Series**:
- MLA-Base: ~25 ms/Token
- EG-MLA-Base: ~26 ms/Token
### Detailed Analysis
1. **MLA-Base**: Consistent latency of ~25 ms/Token
2. **EG-MLA-Base**: Slightly higher latency at ~26 ms/Token
### Key Observations
- Both configurations show nearly identical decoding performance
- EG-MLA-Base incurs minimal latency penalty despite efficiency gains
### Interpretation
The near-parity in decoding latency between configurations suggests that the memory efficiency gains from EG-MLA-Base do not come at the cost of computational speed, making it a favorable choice for resource-constrained environments.
---
## Cross-Quadrant Analysis
### Systemic Trends
1. **Performance-Resource Tradeoff**: Larger embeddings improve accuracy but increase memory requirements
2. **Efficiency Gains**: EG-MLA-Base achieves 91.6% KV cache reduction with only 4% latency increase
3. **Scalability Patterns**: All models show diminishing returns beyond 309M embedded size
### Technical Implications
- The EG-MLA-Base configuration offers optimal balance between performance and efficiency
- MLA-256 provides best accuracy-to-resource ratio for moderate-scale applications
- MHA remains inefficient for modern deployment requirements
### Anomalies
- The slight loss increase at 618M embedded size in EG-MLA-Base (3.07 vs 3.05 at 309M) suggests potential optimization opportunities in intermediate-scale configurations.
</details>
Figure 1: (a) Average accuracy versus training tokens. (b) Validation loss with respect to embedding size used in embedding gating. (c) KV cache elements per token. EG-MLA significantly reduces cache size compared to MHA and MLA. (d) Decoding efficiency of EG-MLA.
Autoregressive decoding remains a critical bottleneck in the deployment of Transformer-based large language models (LLMs), particularly in real-time or resource-constrained scenarios. A major contributor to this inefficiency is the large memory footprint required to store key-value (KV) pairs during inference. In standard Multi-Head Attention (MHA) (Vaswani et al. 2017), each attention head maintains its own set of key and value vectors, leading to quadratic memory scaling with respect to sequence length and model size. While effective in capturing complex dependencies, this design is increasingly impractical for latency-sensitive applications or deployment on edge devices.
To mitigate this challenge, a series of efficient attention mechanisms have been proposed to reduce KV cache requirements while preserving model performance. Notably, Multi-Query Attention (MQA) (Shazeer 2019) and Grouped-Query Attention (GQA) (Ainslie et al. 2023) compress the KV cache by sharing key and value representations across attention heads or head groups. More recently, Multi-head Latent Attention (MLA) (Liu et al. 2024) introduces a shared latent projection space that significantly reduces memory usage while maintaining strong expressiveness, striking a better balance between performance and efficiency.
Building on this line of work, we introduce Embedding-Gated Multi-head Latent Attention (EG-MLA), a novel extension to MLA that further enhances memory efficiency without sacrificing model quality. EG-MLA augments the latent attention mechanism with a token-specific embedding gating module that modulates compressed KV representations in a fine-grained manner. This lightweight gating design introduces minimal additional computation and parameter overhead, yet substantially improves the expressive capacity of the compressed latent space.
Compared to MHA, EG-MLA achieves over 91.6% reduction in KV cache size, and compared to MLA, it provides up to 59.9% additional memory savings while achieving equal or better performance across multiple reasoning benchmarks. Our theoretical analysis reveals that embedding gating implicitly induces a second-order feature space, which enhances token-specific representation capacity. Empirically, we demonstrate that EG-MLA is robust across varying compression ratios, embedding dimensions, and model scales.
Notably, we successfully scale EG-MLA to over 1 billion parameters, confirming its effectiveness not only in compact models but also in large-scale LLMs. Despite aggressive compression, our 1.2B EG-MLA model maintains competitive performance with its MLA counterpart, while using only 40% of the KV cache.
Our contributions can be summarized as follows:
- We propose EG-MLA, an embedding-gated extension to latent attention that enables token-specific modulation of compressed KV representations.
- We demonstrate that EG-MLA achieves superior performance-to-memory trade-offs compared to both MHA and MLA, across diverse tasks and model scales.
- We provide theoretical insights into the gating mechanism and show through comprehensive experiments that EG-MLA supports efficient inference at scale, including in 1B+ parameter regimes.
- We identify that further reducing the KV cache in MLA is inherently difficult, as its latent space is already highly compressed and empirically optimized. To overcome this limitation without degrading model performance, we introduce a novel architectural dimension—embedding gating—which adds token-specific parameters that compensate for the loss of latent capacity. This design preserves model expressiveness and inference efficiency, enabling deeper KV compression with minimal decoding overhead.
Overall, EG-MLA advances the frontier of efficient Transformer architectures, offering a promising path toward scalable, high-performance, and resource-efficient language model deployment.
## Related Works
### KV Cache Compression in Attention Mechanisms
Reducing the key-value (KV) cache size during inference is a fundamental challenge in scaling large language models (LLMs). Traditional multi-head attention (MHA) (Vaswani et al. 2017) stores separate key and value vectors for each attention head, resulting in substantial memory usage. To alleviate this, multi-query attention (MQA) (Shazeer 2019) and grouped-query attention (GQA) (Ainslie et al. 2023) share KV representations across heads or groups, achieving cache reduction with minimal performance degradation.
Multi-head latent attention (MLA) (Liu et al. 2024; Ji et al. 2025) extends this direction by introducing shared latent projections that significantly reduce KV cache size while retaining model expressiveness. Ji et al. (Ji et al. 2025) further adapt MLA to arbitrary Transformer-based LLMs, highlighting its applicability for broad deployment. However, MLA’s reliance on fixed latent projections may limit its adaptability, especially under aggressive compression.
Several recent studies further advance this paradigm. TransMLA (Meng et al. 2025) systematizes the latent attention design for efficient inference, though it does not incorporate token-level adaptability. Deng and Woodland (Deng and Woodland 2025) propose a temporal latent attention mechanism, exploring compression along different axes, which is conceptually related to our gating strategy. Mehta et al. (Mehta et al. 2025) investigate latent attention specifically for small language models, offering evidence that such approaches can maintain quality under limited resource budgets—an area in which EG-MLA also excels.
### Gating Mechanisms in Transformer Architectures
Gating mechanisms have been widely studied to enhance Transformer expressiveness. For example, the Highway Transformer (Chai, Jin, and Hou 2020) employs self-gating to modulate internal information flow, showing that adaptive gating can improve model capacity. However, these approaches typically operate over the full attention space and introduce non-negligible computational overhead.
Our proposed EG-MLA differs by introducing a lightweight embedding gating mechanism that operates in the compressed latent space. This enables token-specific modulation of KV representations with minimal computational and memory cost, effectively enhancing expressiveness without sacrificing efficiency.
### Other Attention Variants
Other attention variants such as Performer (Choromanski et al. 2020), Linformer (Wang et al. 2020), and Longformer (Beltagy, Peters, and Cohan 2020) aim to improve efficiency by approximating the softmax attention kernel. However, these methods primarily target encoder-based architectures or long-context processing, and do not directly address the memory overhead from KV caching during decoder-side autoregressive inference. In contrast, EG-MLA maintains the original attention structure while focusing on minimizing the KV cache footprint during generation. Thus, we consider our comparisons with MLA, MQA, and GQA to be more directly aligned.
### Summary
EG-MLA integrates latent compression and token-specific modulation into a unified architecture. By building on and extending prior work in both KV cache reduction and attention gating, it achieves improved trade-offs between memory usage and model quality. This positions EG-MLA as a practical and scalable solution for efficient LLM inference across a wide range of deployment scenarios.
## Embedding Gated MLA
Simplified Architecture of EG-MLA is shown in Figure 2.
<details>
<summary>x2.png Details</summary>
### Visual Description
## Diagram: EG-MLA Neural Network Architecture
### Overview
The diagram illustrates a technical architecture for an EG-MLA (Efficient Gating-Multi-Head Attention) neural network. It is divided into two main sections:
1. **Left Side (EG-MLA Core)**: Contains components for attention mechanisms, gating, and embeddings.
2. **Right Side (Processing Flow)**: Shows data flow through token and channel mixing layers with normalization.
### Components/Axes
#### Left Side (EG-MLA Core):
- **Linear Layers**:
- Three green "Linear" blocks labeled **Q**, **KV**, and **Emb+** (input embeddings).
- Connected to a **Multi-Head Attention** block (purple).
- **Element-wise Gating**:
- Red block labeled **Element-wise Gating**, receiving input from Multi-Head Attention.
- **Embed Table**:
- Pink block labeled **Emb Table***, connected to **Emb+** via a linear layer.
- Marked with an asterisk (*) indicating it is **precomputed and prefetched during inference**.
- **Offloaded Components**:
- Linear layers for **Q**, **KV**, and **Emb+** are marked with an asterisk (*) as **offloaded during inference**.
#### Right Side (Processing Flow):
- **Input Embedding**:
- Pink block labeled **Input Embedding**, receiving input **idx**.
- **Token Mixing**:
- Orange block labeled **Token Mixing**, connected to **Input Embedding**.
- **Channel Mixing**:
- Blue block labeled **Channel Mixing**, connected to Token Mixing.
- **Add & Norm**:
- Two yellow blocks labeled **Add & Norm**, positioned between Token Mixing and Channel Mixing.
- **xN Multiplier**:
- Label **xN** at the top of the right side, indicating scaling.
### Detailed Analysis
- **Data Flow**:
- Input **idx** → **Input Embedding** → **Token Mixing** → **Add & Norm** → **Channel Mixing** → **Add & Norm** → **xN**.
- Parallel flow from **EG-MLA Core** (Q, KV, Emb+) → **Multi-Head Attention** → **Element-wise Gating** → merged with the right-side flow.
- **Key Connections**:
- **Emb+** (from EG-MLA) is combined with **Input Embedding** (right side) via **Token Mixing**.
- **Multi-Head Attention** and **Element-wise Gating** outputs are integrated into the processing pipeline.
### Key Observations
1. **Efficiency Optimizations**:
- **Precomputed/Prefetched**: The **Emb Table** is precomputed to reduce inference latency.
- **Offloaded**: Linear layers for Q, KV, and Emb+ are offloaded during inference, likely to reduce computational load.
2. **Modular Design**:
- The architecture separates attention/gating logic (left) from token/channel processing (right), suggesting modularity for scalability.
3. **Normalization**:
- **Add & Norm** blocks are used twice, indicating residual connections and layer normalization for stability.
### Interpretation
This architecture combines transformer-like attention mechanisms (Multi-Head Attention) with efficiency optimizations for inference. By offloading certain linear operations and precomputing the embedding table, the model reduces runtime computation. The integration of **Element-wise Gating** suggests a gating mechanism to control information flow, potentially improving performance on specific tasks. The **xN** scaling implies the output is amplified, possibly for task-specific adjustments.
The diagram emphasizes a balance between computational efficiency (via offloading/precomputation) and model complexity (via attention and gating), making it suitable for resource-constrained environments.
</details>
Figure 2: Simplified Architecture of the Embedding Gated Multi-Head Latent Attention (EG-MLA).
### Core Formulas
#### Low-Rank Key-Value Joint Compression
We begin by inheriting the core principle of MLA, namely the low-rank joint compression of keys and values, which serves to reduce the memory footprint of the KV cache:
$$
\displaystylec_t^KV \displaystyle=W^DKVx_t, \displaystylekv_t^C \displaystyle=W^UKVc_t^KV, \tag{1}
$$
where $c_t^KV∈ℝ^d_c$ is the compressed latent vector for keys and values; $W^UKV$ is up-projection matrix for joint keys and values.
#### Embedding Gating
Given the token index $i_t$ , we retrieve its per-layer embedding $e_t$ from a learnable lookup table. This embedding is then passed through an up-projection matrix $W^UE$ to obtain a gating signal $g_t$ , representing a token-specific modulation vector. We apply element-wise modulation to the compressed key-value representation $kv_t^C$ via the Hadamard product with $g_t$ . Finally, a layer normalization is applied to the result, producing the final key and value representation $\widetilde{kv}_t^C$ .
$$
\displaystylee_t \displaystyle=Emb(i_t), \displaystyleg_t \displaystyle=W^UEe_t, \displaystyle\widetilde{kv}_t^C \displaystyle=LN≤ft(kv_t^C\odotg_t\right), \tag{3}
$$
The subsequent computation follows the same procedure as MLA. To provide a comprehensive view of the EG-MLA mechanism, we present the complete set of formulas along with a detailed architectural diagram in Appendix C.
### Rationale behind Embedding Gating
We analyze the embedding gating mechanism by reformulating its core expression, in order to highlight its ability to project inputs into a high-dimensional nonlinear feature space.
In a single layer of EG-MLA, the embedding gating operation can be written as $(W^UKV c_t^KV)\odot(W^UE e_t)$ , which denotes the element-wise multiplication of two linearly projected features.
For ease of analysis, this operation can be abstracted as $(W_1x_1)\odot(W_2x_2)$ , where $x_1∈ℝ^d_1$ and $x_2∈ℝ^d_2$ are two distinct input vectors, and $W_1∈ℝ^d_1× d$ , $W_2∈ℝ^d_2× d$ are their corresponding projection matrices into a shared $d$ -dimensional space.
Focusing on the case of a single output channel, the embedding gating can be further expressed as:
$$
\displaystyle w_1^Tx_1\odot w_2^Tx_2 \displaystyle=≤ft(∑_i=1^d_1w_1^(i)x_1^(i)\right)\odot≤ft(∑_j=1^d_2w_2^(j)x_2^(j)\right) \displaystyle=∑_i=1^d_1∑_j=1^d_2w_1^(i)w_2^(j)x_1^(i)x_2^(j) \tag{6}
$$
As shown in Eq. 8, this formulation expands into $d_1× d_2$ distinct second-order interaction terms between $x_1$ and $x_2$ . Therefore, a single-layer embedding gating operation implicitly maps the input into a nonlinear feature space of dimensionality $d_1× d_2$ , significantly enhancing the representational capacity without introducing any additional computational overhead. This formulation forms the fundamental rationale behind our design choice to incorporate embedding gating. By leveraging its ability to construct an implicit high-dimensional representation, we enable an architecture with a more compact latent space without compromising expressiveness.
## Performance Evaluation
### Experimental Settings
#### Data
We use ClimbMix (Diao et al. 2025) as the training dataset, which consists of 400 billion tokens in total. For evaluation, we adopt several reasoning benchmarks: PIQA (Bisk et al. 2020), ARC-C, ARC-E (Clark et al. 2018), HellaSwag (Zellers et al. 2019), WinoGrande (Sakaguchi et al. 2019), and SIQA (Sap et al. 2019). We optimize the data mixture using the validation sets of PIQA, ARC-E, and HellaSwag, and then evaluate on their respective test sets. All evaluations are performed using the LM-Evaluation Harness (Gao et al. 2024), under the 0-shot setting for all tasks except MMLU, which uses a 5-shot setting (Allal et al. 2025; Cosmopedia Team 2024).
#### Model
We train a series of Transformer decoder-only models using the EG-MLA architecture with next-token prediction as the training objective. The training data consists of 50 billion tokens.
#### Baselines
We compare the proposed EG-MLA architecture against MLA models with equivalent activation parameter counts. More detailed experimental settings and training hyper-parameters can be found in Appendix A.
### Towards Compact KV Cache
Table 1 reveals several noteworthy patterns regarding the impact of embedding gating and KV cache size on model performance. First, the introduction of embedding gating consistently enhances performance across nearly all tasks compared to the baseline MLA model with the same KV cache size, indicating that the proposed mechanism effectively enriches representational capacity without increasing computational burden. Second, EG-MLA models demonstrate strong robustness to aggressive KV compression: even when the KV cache size is reduced from 256 to as low as 64 or 16, the average accuracy remains stable, with only marginal fluctuations across most benchmarks. This suggests that embedding gating allows the model to retain crucial information in a more compact latent space.
| Model Name | Size | piqa | arc-c | arc-e | hella | winog | siqa | mmlu | obqa | boolq | race | truth | lmb.o | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| MLA-Base-kv256 | 120M | 69.48 | 26.45 | 61.36 | 34.11 | 52.33 | 85.40 | 25.12 | 21.40 | 58.10 | 30.91 | 29.82 | 31.63 | 43.84 |
| EG-MLA-Base-kv256 | 125M | 70.29 | 29.27 | 62.80 | 35.82 | 50.67 | 86.40 | 25.22 | 24.40 | 55.26 | 30.62 | 29.20 | 33.63 | 44.46 |
| EG-MLA-Base-kv128 | 122M | 71.49 | 29.18 | 62.75 | 35.57 | 53.43 | 87.80 | 26.71 | 22.20 | 47.74 | 30.24 | 29.09 | 33.48 | 44.14 |
| EG-MLA-Base-kv64 | 120M | 71.60 | 29.61 | 63.89 | 35.31 | 51.14 | 86.20 | 24.19 | 23.00 | 55.05 | 30.43 | 29.21 | 32.35 | 44.33 |
| EG-MLA-Base-kv16 | 119M | 71.16 | 28.50 | 62.08 | 34.48 | 50.75 | 84.70 | 26.21 | 21.40 | 57.13 | 31.29 | 30.58 | 30.95 | 44.10 |
Table 1: Accuracy of different models across 12 benchmark tasks under varying KV cache sizes. All models use a fixed embedding dimension of 256. ”Size” denotes the number of activated parameters. Task abbreviations: LAMBADA (lmb.o), HellaSwag (hella), Winogrande (winog), TruthfulQA (truth).
We also assess EG-MLA on standard language modeling datasets such as WikiText103 (Merity, Keskar, and Socher 2018), LAMBADA (Paperno et al. 2016), and ClimbMix (Diao et al. 2025). As reported in Appendix D, EG-MLA consistently outperforms MLA in perplexity or validation loss under equivalent compression settings. This suggests that embedding gating provides benefits beyond reasoning and applies broadly to language modeling scenarios. These results demonstrate that EG-MLA remains robust even under aggressive compression settings. It is important to emphasize, however, that MLA itself is already near the limit of feasible compression.
### Scaling Law for Embedding Size
Table 2 highlights several key trends as the embedding dimension increases. Most notably, performance consistently improves when moving from extremely small embeddings (e.g., 64) to moderate sizes (e.g., 512), suggesting that larger embedding vectors enable more expressive token-specific modulation through the gating mechanism. However, beyond a certain point—particularly past 1024 dimensions—the performance gains begin to saturate and even exhibit some inconsistency across tasks. For instance, while some benchmarks such as ARC-E and LAMBADA continue to improve steadily, others like MMLU and WinoGrande fluctuate or plateau. This indicates diminishing returns from simply scaling up the embedding size, likely due to a combination of overparameterization and limited benefit from increased modulation granularity. Interestingly, even relatively compact models with 128–256 embedding dimensions already capture most of the benefit, striking a favorable balance between performance and efficiency. These findings suggest that embedding size should be carefully tuned rather than blindly scaled, and that EG-MLA achieves strong performance even under strict parameter budgets.
Notably, performance does not increase monotonically as shown in Figure 1. Beyond the point where the embedding size exceeds the total dimensionality of the key and value vectors, we observe saturation or even slight regressions on some tasks. This suggests that overparameterization beyond the combined key-value capacity yields limited benefits, revealing a current scaling limitation of our method. Exploring architectures that scale more gracefully with increased embedding capacity is an important direction for future work.More comprehensive scaling experiments are provided in Appendix G.
| Model Name | Embed Size | piqa | arc-c | arc-e | hella | winog | siqa | mmlu | obqa | boolq | race | truth | lmb.o | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| EG-MLA-Base-emb64 | 38M | 69.37 | 27.47 | 60.82 | 33.96 | 51.22 | 83.90 | 27.73 | 21.00 | 55.17 | 28.71 | 29.81 | 31.13 | 43.36 |
| EG-MLA-Base-emb128 | 77M | 70.62 | 29.44 | 63.22 | 34.72 | 54.22 | 86.60 | 25.48 | 22.60 | 54.68 | 28.80 | 30.54 | 31.07 | 44.33 |
| EG-MLA-Base-emb256 | 154M | 71.60 | 29.61 | 63.89 | 35.31 | 51.14 | 86.20 | 24.19 | 23.00 | 55.05 | 30.43 | 29.21 | 32.35 | 44.33 |
| EG-MLA-Base-emb512 | 309M | 71.93 | 30.12 | 63.89 | 35.86 | 52.01 | 87.10 | 26.27 | 22.20 | 56.12 | 30.05 | 29.50 | 33.86 | 44.91 |
| EG-MLA-Base-emb1024 | 618M | 71.98 | 31.83 | 64.69 | 36.24 | 50.21 | 86.10 | 26.63 | 23.80 | 58.81 | 29.28 | 28.22 | 32.91 | 45.06 |
| EG-MLA-Base-emb2048 | 1236M | 72.69 | 30.89 | 66.37 | 36.44 | 54.07 | 85.40 | 25.43 | 22.80 | 53.61 | 28.90 | 30.78 | 34.70 | 45.17 |
Table 2: Accuracy results of EG-MLA models across 12 benchmark tasks with varying embedding dimensions. The KV cache dimension is fixed at 64. Task abbreviations: LAMBADA (lmb.o), HellaSwag (hella), WinoGrande (winog), TruthfulQA (truth).
### Ablation on Embedding Gating
The ablation study in Table 3 reveals two critical observations about the design of the embedding gating mechanism.
#### The Critical Role of Layer Normalization
First, the presence of LayerNorm after the gating operation is essential. Its removal leads to a clear degradation in performance, suggesting that normalization is crucial for stabilizing the gated feature distribution and ensuring effective downstream attention computation.
#### Effectiveness of element-wise multiplication
Second, replacing the element-wise multiplication with simple addition results in a noticeable loss increase, although not as severe as removing LayerNorm. This indicates that the element-wise multiplicative interaction plays a more effective role in modulating key-value representations, likely because it enables fine-grained, dimension-wise control over the information flow.
Overall, these findings reinforce the importance of both architectural choices—the normalization step and multiplicative gating—in maximizing the benefits of the embedding-guided modulation.
| Model Name | Size | Tokens | Loss (val) |
| --- | --- | --- | --- |
| MLA-Base-kv256 | 120M | 10B | 3.2156 |
| EG-MLA-Base-kv256 | 125M | 10B | 3.1609 |
| - remove LayerNorm | 125M | 10B | 3.2573 |
| - replace $\odot$ with addition | 125M | 10B | 3.1901 |
Table 3: Ablation on Embedding Gating. ”Size” refers to Activated Parameters
### Scaling to 1B Language Model
To probe into the impact of scaling model size and training data, we train two models, namely MLA-1.2B and EG-MLA-1.2B (with 1.2 billion activated parameters and 600 million offloaded parameters). Both models are trained on 50 billion tokens from the ClimbMix dataset (Diao et al. 2025). MLA-1.2B is built on the Deepseek framework (Liu et al. 2024). EG-MLA-1.2B, which is designed for efficiency, utilizes merely 40.1% of the cache size of MLA-1.2B while still achieving competitive performance.
| Model Name | Size | Embed Size | piqa | arc-c | arc-e | hella | winog | siqa | mmlu | obqa | boolq | race | truth | lmb.o | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| MLA-1.2B | 1188M | 0M | 75.52 | 38.82 | 73.02 | 44.15 | 56.83 | 92.20 | 28.42 | 27.20 | 53.15 | 34.26 | 33.76 | 38.02 | 49.61 |
| EG-MLA-1.2B | 1182M | 618M | 75.41 | 38.74 | 72.85 | 44.88 | 56.83 | 91.00 | 26.92 | 27.80 | 59.45 | 33.97 | 31.27 | 40.25 | 49.95 |
Table 4: Accuracy results of additional models across 12 benchmark tasks for EG-MLA with over 1B parameters. ”Size” denotes the number of activated parameters.
Table 4 demonstrates that EG-MLA-1.2B attains performance on par with MLA-1.2B across a wide array of benchmarks, featuring a remarkable enhancement in average accuracy. This competitive performance, in conjunction with the reduced cache size of EG-MLA-1.2B, emphasizes its capability to generalize proficiently under resource limitations.
In conclusion, EG-MLA-1.2B matches the performance of MLA-1.2B in both base and fine-tuned scenarios, while utilizing merely 40% of the KV cache size of MLA-1.2B. These findings highlight the potential of memory- and compute-efficient architectures for the expansion of large language models, facilitating sustainable and resource-efficient AI development.
## Efficiency Evaluation
### Analytical Comparison of KV Cache Costs
Table 5 presents a theoretical comparison of the KV cache cost per token across different attention mechanisms, including our proposed EG-MLA. Traditional methods such as Multi-Head Attention (MHA) and Grouped-Query Attention (GQA) store separate key and value vectors for each head or group, leading to a cache size that scales with the number of heads $n_h$ or groups $n_g$ . Multi-Query Attention (MQA), while more efficient, sacrifices representational capacity by sharing keys and values across all heads.
In comparison, MLA introduces a shared latent vector to compress key-value representations, significantly reducing the cache size while maintaining reasonable performance. However, EG-MLA further improves this design by applying a token-specific embedding gating mechanism, enabling more expressive modulation of the shared latent without increasing cache overhead. Specifically, EG-MLA stores only a compressed latent vector of dimension $d_c/r$ per token, resulting in a leaner cache footprint than MLA. Despite the more aggressive compression, EG-MLA retains strong representational power thanks to the multiplicative interaction with token embeddings. This analytical comparison highlights EG-MLA’s advantage in offering both compact memory usage and enhanced modeling capacity.
The actual KV cache savings can be found in Appendix F, where we provide a concrete example. In this case, EG-MLA achieves a 59.9% reduction in KV cache compared to MLA, and a substantial 91.6% reduction compared to MHA.
| Attention Mechanism | KV Cache / Token (#Elem) | Capability |
| --- | --- | --- |
| Multi-Head Attention (MHA) | $2n_hd_hl$ | Strong |
| Grouped-Query Attention (GQA) | $2n_gd_hl$ | Moderate |
| Multi-Query Attention (MQA) | $2d_hl$ | Weak |
| Multi-Head Latent Attention (MLA) | $(d_c+d_h^R)l$ | Strong |
| EG-MLA (Ours) | $(d_c/r+d_h^R)l$ | Stronger |
Table 5: Comparison of the number of KV cache elements per token across different attention mechanisms. $n_h$ denotes the number of attention heads, $d_h$ the dimension per attention head, $l$ the number of layers, $n_g$ the number of groups in GQA, and $d_c$ and $d_h^R$ represent the latent space dimension and the KV cache dimension used for RoPE in MLA, respectively. For EG-MLA, a compressed latent vector of dimension $d_c/r$ is stored per token, where $r$ is the compression factor applied to the latent representation. This results in a significantly smaller KV cache while preserving strong representational capacity through embedding gating.
<details>
<summary>x3.png Details</summary>
### Visual Description
## Flowchart: Apply Embedding to kv
### Overview
The diagram illustrates a sequential data processing pipeline for applying embeddings to a knowledge vector (kv). It depicts a structured workflow with six key components connected by directional arrows, representing the flow of data transformations.
### Components/Axes
1. **Input**:
- **Emb** (embedding layer) - Yellow box on the far left
2. **Processing Steps**:
- **e1** (embedding output 1) - Light blue box
- **W** (weight matrix) - Peach box (appears twice)
- **Gather&Cat** (gather and concatenate) - Light blue box
- **Norm** (normalization) - Pink box
3. **Output**:
- Final normalized vector (arrow exiting right side)
### Detailed Analysis
1. **Flow Sequence**:
- **Emb** → **e1** → **W** → **Gather&Cat** → **Norm**
- **Emb** → **e2** → **W** → **Gather&Cat** → **Norm**
2. **Key Operations**:
- **W**: Weight matrix multiplication applied to both e1 and e2
- **Gather&Cat**: Combines transformed embeddings (e1*W and e2*W)
- **Norm**: Final normalization step
### Key Observations
- Dual-path processing: Two parallel embedding streams (e1 and e2) converge after weight matrix application
- Normalization is the final processing step before output
- No numerical values or quantitative metrics present in the diagram
### Interpretation
This diagram represents a typical neural network architecture component for knowledge vector processing. The workflow suggests:
1. **Embedding Layer**: Converts input features into dense vector representations (e1 and e2)
2. **Weight Transformation**: Applies learned parameters (W) to embeddings
3. **Feature Fusion**: Combines transformed embeddings through concatenation
4. **Normalization**: Ensures output vector has unit magnitude, crucial for stable training dynamics
The architecture implies a Siamese-like network design where two embedding paths are processed identically before fusion, commonly used in contrastive learning tasks. The explicit normalization step indicates attention to gradient stability and output scale consistency.
</details>
Figure 3: Dimension-wise sharding pipeline.
### Empirical Comparison of Efficiency
We evaluate the inference efficiency of EG-MLA on a single NVIDIA A800 80 GB GPU. As shown in Figure 4, we report Prefill latency, Decoding latency and tokens-per-second under varying prompt lengths and batch sizes. The MLA model contains 120M parameters in total, whereas EG-MLA contains 274M, of which 154M are embed parameters.EG-MLA-A is an inference-optimized variant of EG-MLA in which the projection matrices are pre-computed offline during model loading; during inference, matrix multiplication is thus eliminated and only a table lookup is performed, further reducing inference latency.Despite the extra embed parameters introduced, EG-MLA maintains inference performance comparable to or better than that of MLA. Concretely, when the batch size exceeds 8, EG-MLA reduces the average pre-fill latency by 5.6 ms compared with MLA while adding only 0.91 ms to the average decoding latency. When the batch size reaches 32, EG-MLA-A narrows the average decoding latency gap with MLA to merely 1 ms, underscoring the efficiency of our approach for large-scale parallel workloads.
<details>
<summary>x4.png Details</summary>
### Visual Description
## Bar Chart: Performance Comparison Across Models and Batch Sizes
### Overview
The image contains nine grouped bar charts comparing three machine learning models (MLA, EG-MLA, EG-MLA-A) across three metrics (Prefill Time, Decoding Time, Tokens/second) for different prompt lengths (1024, 2048, 4096 tokens). Each metric is visualized across batch sizes (1, 4, 8, 16, 32) with distinct color coding for model types.
### Components/Axes
- **X-axis**: Batch Size (1, 4, 8, 16, 32)
- **Y-axes**:
- Prefill Time (ms): Logarithmic scale (0, 100, 500, 1000, 2000)
- Decoding Time (ms): Linear scale (0, 10, 20, 30, 40, 50, 60)
- Tokens/second: Linear scale (0, 20, 40, 60, 80)
- **Legend**:
- Black: MLA
- Light Blue: EG-MLA
- Dark Blue: EG-MLA-A
- **Chart Groups**:
- Top row: `prompt_len_1024`
- Middle row: `prompt_len_2048`
- Bottom row: `prompt_len_4096`
### Detailed Analysis
#### Prefill Time
- **prompt_len_1024**:
- MLA: 29.53 → 53.0 ms (↑82.9%)
- EG-MLA: 27.38 → 52.7 ms (↑92.5%)
- EG-MLA-A: 27.38 → 53.0 ms (↑93.6%)
- **prompt_len_2048**:
- MLA: 50.86 → 219.9 ms (↑332.3%)
- EG-MLA: 159 → 219.9 ms (↑37.7%)
- EG-MLA-A: 159 → 219.9 ms (↑37.7%)
- **prompt_len_4096**:
- MLA: 97.89 → 296.6 ms (↑202.7%)
- EG-MLA: 364 → 296.6 ms (↓18.5%)
- EG-MLA-A: 364 → 296.6 ms (↓18.5%)
#### Decoding Time
- **prompt_len_1024**:
- MLA: 13.6 → 21.3 ms (↑56.6%)
- EG-MLA: 14.6 → 19.7 ms (↑34.9%)
- EG-MLA-A: 14.6 → 17.7 ms (↑21.2%)
- **prompt_len_2048**:
- MLA: 15.6 → 25.6 ms (↑64.1%)
- EG-MLA: 16.6 → 27.7 ms (↑67.2%)
- EG-MLA-A: 16.6 → 29.0 ms (↑74.7%)
- **prompt_len_4096**:
- MLA: 20.0 → 40.4 ms (↑102%)
- EG-MLA: 21.2 → 40.3 ms (↑89.6%)
- EG-MLA-A: 21.2 → 40.3 ms (↑89.6%)
#### Tokens/second
- **prompt_len_1024**:
- MLA: 79.7 → 136.1 (↑70.7%)
- EG-MLA: 364 → 136.1 (↓62.6%)
- EG-MLA-A: 364 → 136.1 (↓62.6%)
- **prompt_len_2048**:
- MLA: 49.5 → 168.3 (↑240.2%)
- EG-MLA: 273.7 → 168.3 (↓38.5%)
- EG-MLA-A: 273.7 → 168.3 (↓38.5%)
- **prompt_len_4096**:
- MLA: 49.5 → 273.7 (↑452.3%)
- EG-MLA: 40.4 → 273.7 (↓35.6%)
- EG-MLA-A: 40.4 → 273.7 (↓35.6%)
### Key Observations
1. **Batch Size Impact**:
- Prefill Time increases non-linearly with batch size for all models.
- Decoding Time shows linear scaling with batch size.
- Tokens/second decreases with larger batches, but MLA maintains higher throughput.
2. **Model Performance**:
- **EG-MLA-A** dominates in Decoding Time and Tokens/second across all prompts.
- **MLA** has the highest Prefill Time but best throughput for small batches (e.g., 16 tokens/second at batch 32 for 4096-length prompts).
- **EG-MLA** underperforms in Prefill Time for large prompts (4096) but matches MLA in decoding efficiency.
3. **Anomalies**:
- EG-MLA-A shows inconsistent Prefill Time trends (improves for 2048, worsens for 4096).
- MLA's Prefill Time spikes disproportionately for 4096-length prompts.
### Interpretation
The data reveals a trade-off between Prefill Time and decoding efficiency. EG-MLA-A optimizes Decoding Time and throughput at the cost of increased Prefill Time for large prompts. MLA, while slower in Prefill, maintains competitive throughput for smaller batches. The inconsistent Prefill Time behavior of EG-MLA-A suggests potential optimization opportunities for large-context tasks. These findings highlight the importance of model selection based on application-specific requirements (e.g., real-time vs. batch processing).
</details>
Figure 4: Prefill and Decoding times and Tokens per seconds for MLA, EG-MLA and EG-MLA-A across configurations on an NVIDIA A800 80GB GPU. EG-MLA-A is an inference-optimized variant of EG-MLA that pre-computes all embeddings at model-loading time and pre-fetch during inference.
| Method | Node | Peak Mem. | Throughput | Avg. Step Time | Comm. Time |
| --- | --- | --- | --- | --- | --- |
| vanilla | 2 | 13950 | 42528 | 385.25 | – |
| dimension-wise | 2 | 12233 | 86743 | 188.88 | 16.25 |
| vanilla | 4 | 13950 | 80114 | 409.02 | – |
| dimension-wise | 4 | 11759 | 175031 | 187.21 | 24.38 |
Table 6: Performance comparison between vanilla and dimension-wise parallelism under varying node counts. Memory is reported in megabytes (MB) and throughput in tokens per second (tokens/s), and all timing metrics in milliseconds (ms).
### Embedding Parallelism
Embedding gating introduces additional parameters into each attention layer, significantly increasing GPU memory usage during training. To mitigate peak memory footprint on individual devices, we adopt a dimension-wise sharding strategy (Shoeybi et al. 2020), which partitions the key–value embedding uniformly along the embedding dimension across $W$ parallel workers. The per-rank embedding dimension is given by:
$$
\displaystyled_local \displaystyle=\frac{d_kv}{W}, \tag{9}
$$
where $d_kv$ is the original KV embedding dimension.
Each worker independently computes its local embedding and corresponding projection:
$$
\displaystylee_i \displaystyle=Emb(i_t), \displaystyleg_t^(i) \displaystyle=W_i^UEe_i, \tag{10}
$$
and the final gating vector is reconstructed via concatenation.
This approach yields near-linear memory scaling. Embedding lookups and projections are fully parallelized across workers, with only a single all-gather operation needed during final aggregation, incurring minimal communication overhead. A simplified pipeline is illustrated in Figure 3.
As shown in Table 6, dimension-wise parallelism with two nodes reduces peak memory by 12.3%, shortens per-step time by 51%, and increases throughput by 104% relative to the baseline. Scaling to four nodes yields further reductions in both peak memory and per-step time, together with an additional throughput gain.
Why Dimension-wise Sharding?
Tensor parallelism distributes matrix computations across multiple GPUs to reduce memory usage and accelerate training, but introduces inter-device communication. Compared to alternatives like layer-wise sharding, we choose dimension-wise partitioning for the following reasons:
- Load Balancing: Evenly splits the embedding matrix, ensuring uniform computational load across devices.
- No Serial Dependencies: Enables full parallelism by avoiding sequential computation across layers.
- Memory Efficiency: Embedding layers are memory-intensive; partitioning them maximizes GPU memory bandwidth and cache utilization.
<details>
<summary>x5.png Details</summary>
### Visual Description
## Line Graphs: PPL vs Embedding Dimension and PPL vs Layer Group Size
### Overview
The image contains two line graphs comparing Perplexity (PPL) metrics across different model configurations. Graph (a) examines PPL relative to embedding dimensions, while graph (b) analyzes PPL against layer group sizes. Both graphs include baseline reference lines for comparison.
### Components/Axes
**Graph (a): PPL vs Embedding Dimension**
- **X-axis**: Embedding Dimension (values: 64, 128, 256, 512, 1024, 2048)
- **Y-axis**: Avg. PPL (range: 21–28)
- **Legend**:
- Blue line: "EG-MLA-emb"
- Red dashed line: "MLA-kv256"
**Graph (b): PPL vs Layer Group Size**
- **X-axis**: Layer Group Size (values: 0, 2, 4, 6)
- **Y-axis**: Avg. PPL (range: 11.5–13.5)
- **Legend**:
- Green line: "EG-MLA-Large-ctx2048-kv256emb512"
- Red dashed line: "MLA-Large-ctx2048-kv512"
### Detailed Analysis
**Graph (a) Trends**:
1. **EG-MLA-emb (blue line)**:
- Starts at ~27.8 PPL at 64 dimensions
- Sharp decline to ~25 PPL at 128 dimensions
- Further drop to ~23 PPL at 256 dimensions
- Slight increase to ~20.7 PPL at 512 dimensions
- Stabilizes at ~21 PPL at 1024 and 2048 dimensions
2. **MLA-kv256 (red dashed line)**: Constant at 27 PPL across all dimensions
**Graph (b) Trends**:
1. **EG-MLA-Large-ctx2048-kv256emb512 (green line)**:
- Starts at ~11.6 PPL at layer group size 0
- Increases to ~11.9 PPL at size 2
- Rises to ~12.5 PPL at size 4
- Peaks at ~13.5 PPL at size 6
2. **MLA-Large-ctx2048-kv512 (red dashed line)**: Constant at 12 PPL across all sizes
### Key Observations
1. **Embedding Dimension Impact**:
- PPL decreases significantly with increasing dimensions up to 256
- Diminishing returns observed beyond 256 dimensions
- EG-MLA-emb consistently outperforms MLA-kv256 baseline
2. **Layer Group Size Impact**:
- PPL increases linearly with larger layer groups
- EG-MLA-Large model shows 16.7% PPL increase from size 0 to 6
- MLA-Large baseline remains stable at 12 PPL
### Interpretation
The data suggests:
1. **Optimal Embedding Dimensions**:
- 256 dimensions appear to be the "sweet spot" for EG-MLA-emb, balancing PPL reduction and computational efficiency
- Beyond 256 dimensions, PPL improvements plateau while resource costs likely increase
2. **Layer Group Size Tradeoffs**:
- Larger layer groups (6 vs 0) result in 16.7% higher PPL for EG-MLA-Large
- This indicates potential overfitting or increased model complexity costs
- The stable MLA-Large baseline suggests architectural differences mitigate PPL growth
3. **Model Architecture Insights**:
- EG-MLA variants consistently outperform MLA baselines
- The "kv256" and "kv512" configurations show different scaling behaviors
- The "ctx2048" context length appears critical for maintaining PPL performance
These findings highlight the importance of balancing model complexity with PPL performance, suggesting that optimal configurations exist beyond simple "bigger is better" approaches.
</details>
Figure 5: (a) Average Perplexity(PPL) versus Embedding Dimension. (b) Average Perplexity(PPL) versus Layer Group Size in Inter. The PPL reported in both (a) and (b) are averaged over the ClimbMix-val, Wiki, and LAMBADA benchmarks.
### Toward Greater Efficiency
To compress model parameters further while preserving performance, we explore two orthogonal directions. First, as shown in Table 5, the proposed embedding-gating mechanism yields a significantly smaller KV-cache than competing techniques. Second, we investigate inter-layer parameter sharing: the first half of the KV-projection matrices is reused across neighboring layers, with a single hyper-parameter that controls the width of the sharing window. Setting this hyper-parameter to 2, for example, theoretically halves the parameter count of the first half of the KV projections and simultaneously reuses the associated intermediate activations, cutting both memory footprint and computation. Detailed ablations are provided in Appendix H.
Figure 5 reveals that perplexity decreases monotonically with embedding dimension (emb) until it plateaus beyond emb = 512, indicating the current capacity limit of our embedding gating mechanism; surpassing this limit will be the focus of future work. Figure 3(b) shows that a layer group size of 2 or 3 strikes the best trade-off. Consequently, we recommend emb = 512 and layer group size = 2 for an optimal balance between performance and storage.
## Conclusion
In this work, we introduced Embedding-Gated Multi-head Latent Attention (EG-MLA), a novel extension of MLA designed to achieve more compact and expressive key-value (KV) representations. By integrating a lightweight embedding gating mechanism, EG-MLA effectively compresses the KV cache without sacrificing model performance. Through extensive experiments across multiple reasoning benchmarks, we demonstrated that EG-MLA consistently matches or surpasses MLA while significantly reducing memory usage, achieving up to a 59.9% reduction in KV cache compared to MLA and over 91.6% relative to MHA. We further analyzed the impact of embedding size, layer normalization, and multiplicative gating, highlighting the critical design choices that underpin EG-MLA’s effectiveness. Although MLA already represents a compact and efficient KV caching solution, our work highlights that meaningful further compression is only achievable by innovating along other architectural axes. EG-MLA achieves this by introducing a new token-specific embedding gating dimension that improves expressiveness with minimal impact on decoding latency, thereby enabling deeper compression without performance degradation. Scaling experiments on 1B-parameter models confirmed that EG-MLA retains competitive accuracy under stringent resource constraints, making it a promising candidate for efficient large language model deployment. Additionally, our theoretical and empirical efficiency analyses reveal that embedding gating enables a favorable trade-off between expressiveness and memory footprint. Future work will explore adaptive gating strategies, integration with sparsity-based attention mechanisms, and broader applications of EG-MLA in multi-modal and instruction-tuned models. We believe EG-MLA provides a foundation for advancing memory-efficient and high-performance language model architectures.
## References
- Ainslie et al. (2023) Ainslie, J.; Lee-Thorp, J.; De Jong, M.; Zemlyanskiy, Y.; Lebrón, F.; and Sanghai, S. 2023. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint arXiv:2305.13245.
- Allal et al. (2025) Allal, L. B.; Lozhkov, A.; Bakouch, E.; Bl’azquez, G. M.; Penedo, G.; Tunstall, L.; Marafioti, A.; Kydl’ivcek, H.; Lajar’in, A. P.; Srivastav, V.; Lochner, J.; Fahlgren, C.; Nguyen, X.-S.; Fourrier, C.; Burtenshaw, B.; Larcher, H.; Zhao, H.; Zakka, C.; Morlon, M.; Raffel, C.; von Werra, L.; and Wolf, T. 2025. SmolLM2: When Smol Goes Big - Data-Centric Training of a Small Language Model. ArXiv, abs/2502.02737.
- Beltagy, Peters, and Cohan (2020) Beltagy, I.; Peters, M. E.; and Cohan, A. 2020. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150.
- Bisk et al. (2020) Bisk, Y.; Zellers, R.; Gao, J.; and Choi, Y. 2020. PIQA: Reasoning about Physical Commonsense in Natural Language. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, 7432–7439.
- Chai, Jin, and Hou (2020) Chai, Y.; Jin, S.; and Hou, X. 2020. Highway transformer: Self-gating enhanced self-attentive networks. arXiv preprint arXiv:2004.08178.
- Choromanski et al. (2020) Choromanski, K.; Likhosherstov, V.; Dohan, D.; Song, X.; Gane, A.; Sarlos, T.; Hawkins, P.; Davis, J.; Mohiuddin, A.; Kaiser, L.; et al. 2020. Rethinking attention with performers. arXiv preprint arXiv:2009.14794.
- Clark et al. (2018) Clark, P.; Cowhey, I.; Etzioni, O.; Khot, T.; Sabharwal, A.; Schoenick, C.; and Tafjord, O. 2018. Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge. ArXiv, abs/1803.05457.
- Cosmopedia Team (2024) Cosmopedia Team. 2024. Cosmopedia/evaluation at main · huggingface/cosmopedia. https://huggingface.co/Cosmopedia/evaluation.
- Deng and Woodland (2025) Deng, K.; and Woodland, P. C. 2025. Multi-head Temporal Latent Attention. arXiv preprint arXiv:2505.13544.
- Diao et al. (2025) Diao, S.; Yang, Y.; Fu, Y.; Dong, X.; Su, D.; Kliegl, M.; Chen, Z.; Belcák, P.; Suhara, Y.; Yin, H.; Patwary, M.; Lin, Y. C.; Kautz, J.; and Molchanov, P. 2025. CLIMB: CLustering-based Iterative Data Mixture Bootstrapping for Language Model Pre-training. ArXiv, abs/2504.13161.
- Gao et al. (2024) Gao, L.; Tow, J.; Abbasi, B.; Biderman, S.; Black, S.; DiPofi, A.; Foster, C.; Golding, L.; Hsu, J.; Le Noac’h, A.; Li, H.; McDonell, K.; Muennighoff, N.; Ociepa, C.; Phang, J.; Reynolds, L.; Schoelkopf, H.; Skowron, A.; Sutawika, L.; Tang, E.; Thite, A.; Wang, B.; Wang, K.; and Zou, A. 2024. The Language Model Evaluation Harness.
- Ji et al. (2025) Ji, T.; Guo, B.; Wu, Y.; Guo, Q.; Shen, L.; Chen, Z.; Qiu, X.; Zhang, Q.; and Gui, T. 2025. Towards Economical Inference: Enabling DeepSeek’s Multi-Head Latent Attention in Any Transformer-based LLMs. arXiv preprint arXiv:2502.14837.
- Liu et al. (2024) Liu, A.; Feng, B.; Wang, B.; Wang, B.; Liu, B.; Zhao, C.; Dengr, C.; Ruan, C.; Dai, D.; Guo, D.; et al. 2024. Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model. arXiv preprint arXiv:2405.04434.
- Mehta et al. (2025) Mehta, S.; Dandekar, R.; Dandekar, R.; and Panat, S. 2025. Latent Multi-Head Attention for Small Language Models. arXiv preprint arXiv:2506.09342.
- Meng et al. (2025) Meng, F.; Tang, P.; Tang, X.; Yao, Z.; Sun, X.; and Zhang, M. 2025. Transmla: Multi-head latent attention is all you need. arXiv preprint arXiv:2502.07864.
- Merity, Keskar, and Socher (2018) Merity, S.; Keskar, N. S.; and Socher, R. 2018. An analysis of neural language modeling at multiple scales. arXiv preprint arXiv:1803.08240.
- Paperno et al. (2016) Paperno, D.; Kruszewski, G.; Lazaridou, A.; Pham, Q. N.; Bernardi, R.; Pezzelle, S.; Baroni, M.; Boleda, G.; and Fernández, R. 2016. The LAMBADA dataset: Word prediction requiring a broad discourse context. arXiv preprint arXiv:1606.06031.
- Sakaguchi et al. (2019) Sakaguchi, K.; Bras, R. L.; Bhagavatula, C.; and Choi, Y. 2019. WinoGrande. Communications of the ACM, 64: 99 – 106.
- Sap et al. (2019) Sap, M.; Rashkin, H.; Chen, D.; Le Bras, R.; and Choi, Y. 2019. Social IQa: Commonsense Reasoning about Social Interactions. In Inui, K.; Jiang, J.; Ng, V.; and Wan, X., eds., Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), 4463–4473. Hong Kong, China: Association for Computational Linguistics.
- Shazeer (2019) Shazeer, N. 2019. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.02150.
- Shoeybi et al. (2020) Shoeybi, M.; Patwary, M.; Puri, R.; LeGresley, P.; Casper, J.; and Catanzaro, B. 2020. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv:1909.08053.
- Vaswani et al. (2017) Vaswani, A.; Shazeer, N.; Parmar, N.; Uszkoreit, J.; Jones, L.; Gomez, A. N.; Kaiser, Ł.; and Polosukhin, I. 2017. Attention is all you need. Advances in neural information processing systems, 30.
- Wang et al. (2020) Wang, S.; Li, B. Z.; Khabsa, M.; Fang, H.; and Ma, H. 2020. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768.
- Zellers et al. (2019) Zellers, R.; Holtzman, A.; Bisk, Y.; Farhadi, A.; and Choi, Y. 2019. HellaSwag: Can a Machine Really Finish Your Sentence? In Annual Meeting of the Association for Computational Linguistics.
## A. Hyper-parameters
| Model Name | Model Size | Embed Size | n_layer / n_emb / kv_emb_dim | q_lora_rank / kv_lora_rank | qk_head_dim / v_head_dim | Batch Size | Learning Rate |
| --- | --- | --- | --- | --- | --- | --- | --- |
| MLA-Base-kv256 | 120M | 0M | 12/768/0 | 0/256 | 64+64 / 64 | 0.39M | $6.0× 10^-4$ |
| MLA-Large-kv512 | 645M | 0M | 12/2048/0 | 0/512 | 128+64 / 128 | 0.39M | $2.5× 10^-4$ |
| MLA-XLarge-kv512 | 1188M | 0M | 24/2048/0 | 0/512 | 128+64 / 128 | 0.39M | $2.0× 10^-4$ |
| EG-MLA-Base-kv256 | 279M | 154M | 12/768/256 | 0/256 | 64+64 / 64 | 0.39M | $6.0× 10^-4$ |
| EG-MLA-Base-kv128 | 276M | 154M | 12/768/256 | 0/128 | 64+64 / 64 | 0.39M | $6.0× 10^-4$ |
| EG-MLA-Base-kv64 | 274M | 154M | 12/768/256 | 0/64 | 64+64 / 64 | 0.39M | $6.0× 10^-4$ |
| EG-MLA-Base-kv16 | 273M | 154M | 12/768/256 | 0/16 | 64+64 / 64 | 0.39M | $6.0× 10^-4$ |
| EG-MLA-Base-emb64 | 155M | 38M | 12/768/64 | 0/64 | 64+64 / 64 | 0.39M | $6.0× 10^-4$ |
| EG-MLA-Base-emb128 | 195M | 77M | 12/768/128 | 0/64 | 64+64 / 64 | 0.39M | $6.0× 10^-4$ |
| EG-MLA-Base-emb256 | 274M | 154M | 12/768/256 | 0/64 | 64+64 / 64 | 0.39M | $6.0× 10^-4$ |
| EG-MLA-Base-emb512 | 433M | 309M | 12/768/512 | 0/64 | 64+64 / 64 | 0.39M | $6.0× 10^-4$ |
| EG-MLA-Base-emb1024 | 752M | 618M | 12/768/1024 | 0/64 | 64+64 / 64 | 0.39M | $6.0× 10^-4$ |
| EG-MLA-Base-emb2048 | 1389M | 1236M | 12/768/2048 | 0/64 | 64+64 / 64 | 0.39M | $6.0× 10^-4$ |
| EG-MLA-XLarge-kv128emb512 | 1800M | 618M | 24/2048/512 | 0/128 | 128+64 / 128 | 0.39M | $2.0× 10^-4$ |
| EG-MLA-Large-kv128emb512 | 274M | 154M | 12/2048/512 | 0/128 | 128+64 / 128 | 0.39M | $2.5× 10^-4$ |
| EG-MLA-Large-kv256emb512 | 274M | 154M | 12/768/512 | 0/256 | 64+64 / 64 | 0.39M | $2.5× 10^-4$ |
Table 7: Sizes, architectures, and learning hyper-parameters (batch size in tokens and learning rate) of the models which we trained. All models were trained for a total of 50 billion tokens.
## B. Why Reducing KV Cache Size Matters
### Why is it so important to reduce the size of the KV cache?
It is well known that inference for large language models (LLMs) is typically performed on GPUs, where memory is a scarce and valuable resource. On a single GPU, memory is primarily consumed by two components: (1) model parameters and activation values required for forward computation, which are relatively fixed once the model size is determined; and (2) the KV cache, which stores intermediate key and value tensors for autoregressive decoding. Unlike the model parameters, the KV cache grows dynamically with the input context length and can dominate total memory usage as the sequence becomes longer. In extreme cases, the KV cache alone can exceed the memory capacity of a single GPU or even an entire 8-GPU server.
A fundamental principle of LLM deployment is: if the model can run on a single GPU, do not use multiple GPUs; if it can run on a single server, do not distribute it across machines. This is due to the hardware hierarchy: intra-GPU communication bandwidth is much higher than inter-GPU, which in turn is much higher than inter-node bandwidth. As a result, the more devices involved in inference, the more likely it is to suffer from communication bottlenecks—a classic manifestation of the bottleneck effect. In fact, even the H100 GPU, with up to 3 TB/s of internal SRAM-to-HBM bandwidth, can still experience bandwidth constraints for short-context inference. These issues are exacerbated when spanning multiple GPUs or nodes with slower interconnects.
Therefore, reducing KV cache size enables inference over longer contexts using fewer devices, or allows larger batch sizes under the same memory budget. This results in either faster inference or higher overall throughput. Ultimately, all of these benefits translate to lower inference cost, which is a key concern in the practical deployment of LLMs.
## C. Detailed Architecture and Full Formulas
Detailed Architecture is shown in Figure 6.
The full computation pipeline begins by projecting the input into a low-rank latent space, followed by an up-projection to the original dimensionality. Embedding gating is then applied to produce the key and value representations for attention. These representations are subsequently split, with Partial RoPE individually applied to the keys and values, before being fed into a multi-head attention (MHA) module for final computation. Full formulas are shown as follows:
$$
\displaystylec_t^KV \displaystyle=W^DKVx_t, \displaystylekv_t^C \displaystyle=W^UKVc_t^KV, \displaystylee_t \displaystyle=Emb(i_t), \displaystyleg_t \displaystyle=W^UEe_t, \displaystyle\widetilde{kv}_t^C \displaystyle=LN≤ft(kv_t^C\odotg_t\right), \displaystylek_t^C,v_t^C \displaystyle=split(\widetilde{kv}_t^C) \displaystylec_t^Q \displaystyle=W^DQx_t, \displaystyleq_t^C \displaystyle=W^UQc_t^Q, \displaystyleq_t^R \displaystyle=({W^QR}c_t^Q), \displaystylek_t^R \displaystyle=({W^KR}x_t), \displaystyleq_t,i \displaystyle=[q_t,i^C;q_t,i^R], \displaystylek_t,i \displaystyle=[k_t,i^C;k_t^R], \displaystyleu_t,i \displaystyle=∑_j=1^t_j(\frac{q_t,i^Tk_j,i}{√{d_h+d_h^R}})v_j,i^C, \displaystyleo_t \displaystyle=W^O[u_t,1;u_t,2;...;u_t,n_{h}], \tag{1}
$$
<details>
<summary>x6.png Details</summary>
### Visual Description
## Diagram: Embedding Gated Multi-Head Latent Attention (EG-MLA)
### Overview
The diagram illustrates the architecture of the EG-MLA model, a neural network component that integrates multi-head attention with latent vector processing and key-value caching. The flow begins with an input sequence, processes it through multi-head attention mechanisms, and produces an output sequence. Key components include concatenation, replication, RoPE (Rotary Positional Encoding), element-wise gating, normalization, and projection steps.
### Components/Axes
- **Input**: Labeled as `x_t` (blue rectangle at the bottom-left).
- **Multi-Head Attention**: Central block with sub-components:
- **Concatenate**: Splits into `{q_C}` (green) and `{q_R}` (green).
- **Replicate**: Copies `{k_R}` (white).
- **Apply RoPE**: Processes `{k_R}` (white).
- **Element-wise Gating & LayNorm**: Combines `{k_C}` (yellow) and `{v_C}` (pink).
- **RMS Norm & Projection**: Processes normalized values.
- **Projection**: Final step before output.
- **Output**: Labeled as `o_t` (blue rectangle at the top-right).
- **Latent Vectors**:
- `c_t^Q` (green) and `c_t^KV` (red) are latent representations.
- **Look Up Table**: Pink block with `embed e_t` and `index i_t`.
- **KV Cache**: Red dots labeled `*KV Cache`.
### Detailed Analysis
1. **Input Processing**:
- Input `x_t` is split into concatenated queries `{q_C}` and `{q_R}`.
- Queries are further split into `{q_C}` and `{q_R}` for separate processing.
2. **Multi-Head Attention Flow**:
- `{k_R}` is replicated and processed via RoPE.
- `{k_C}` and `{v_C}` undergo element-wise gating and LayNorm.
- Normalized values are projected and combined with the KV Cache.
3. **Output Generation**:
- The final output `o_t` is derived from the projected values after multi-head attention.
### Key Observations
- **Latent Attention Integration**: The model combines latent vectors (`c_t^Q`, `c_t^KV`) with multi-head attention, suggesting a hybrid approach for sequence modeling.
- **KV Cache**: The red dots indicate a memory mechanism for storing key-value pairs, likely for efficient autoregressive generation.
- **RoPE and Normalization**: The use of RoPE and RMS normalization implies handling of positional information and stability in training.
### Interpretation
The EG-MLA model appears to enhance traditional multi-head attention by incorporating latent vectors and a key-value cache, potentially improving efficiency and context retention. The integration of RoPE and LayNorm suggests a focus on positional awareness and normalization, which are critical for transformer-based models. The KV Cache implies that the model retains historical information, enabling it to process sequences more effectively, especially in autoregressive tasks. This architecture could be applied to tasks like language modeling or sequence-to-sequence generation, where maintaining context and efficiency are paramount.
</details>
Figure 6: Detailed Architecture of the Embedding Gated Multi-Head Latent Attention (EG-MLA).
## D. More Experimental Results
Here we present more experimental data, covering additional dimensions, a broader range of metrics, and more test datasets.
### Scaling Embedding Size
| Model Name | Embed Size | Tokens | ClimbMix-val | wiki | lambada | Avg. |
| --- | --- | --- | --- | --- | --- | --- |
| EG-MLA-emb64 | 38M | 50B | 15.47 | 15.49 | 50.38 | 27.78 |
| EG-MLA-emb128 | 77M | 50B | 15.04 | 15.04 | 45.11 | 25.06 |
| EG-MLA-emb256 | 154M | 50B | 15.32 | 14.58 | 38.94 | 22.95 |
| EG-MLA-emb512 | 309M | 50B | 14.61 | 14.02 | 32.89 | 20.51 |
| EG-MLA-emb1024 | 618M | 50B | 13.91 | 13.72 | 35.41 | 21.01 |
| EG-MLA-emb2048 | 1236M | 50B | 13.81 | 13.62 | 34.25 | 20.56 |
Table 8: Perplexity results of EG-MLA models with scaling embedding size. The KV cache is fixed at 64.
### Scaling KV Cache
| Model Name | Size | Tokens | ClimbMix-val | wiki | lambada | Avg. |
| --- | --- | --- | --- | --- | --- | --- |
| MLA-kv256 | 120M | 50B | 15.45 | 15.63 | 49.88 | 26.99 |
| EG-MLA-kv256 | 125M | 50B | 14.32 | 14.26 | 37.12 | 21.90 |
| EG-MLA-kv128 | 122M | 50B | 14.49 | 14.34 | 35.73 | 21.52 |
| EG-MLA-kv64 | 120M | 50B | 15.32 | 14.58 | 38.94 | 22.95 |
| EG-MLA-kv16 | 119M | 50B | 15.13 | 15.08 | 43.76 | 24.66 |
Table 9: Perplexity results of EG-MLA models across 12 tasks with scaling KV cache. The embedding size is fixed at 256. ”Size” denotes the number of activated parameters.
## E. PyTorch Code
We provide the PyTorch code here to facilitate readers’ understanding of our method.
⬇
1 class EmbeddingGatedMLA (nn. Module):
2 def __init__ (self, config):
3 super (). __init__ ()
4 self. n_embd = config. n_embd
5 self. n_head = config. n_head
6 self. qk_nope_head_dim = config. qk_nope_head_dim
7 self. qk_rope_head_dim = config. qk_rope_head_dim
8 self. qk_head_dim = config. qk_nope_head_dim + config. qk_rope_head_dim
9 self. v_head_dim = config. v_head_dim
10 self. q_lora_rank = config. q_lora_rank
11 self. kv_lora_rank = config. kv_lora_rank
12 assert config. n_embd % config. n_head == 0
13
14 if self. q_lora_rank == 0:
15 self. wq = nn. Linear (config. n_embd, config. n_head * self. qk_head_dim, bias = False)
16 else:
17 self. wq_a = nn. Linear (config. n_embd, config. q_lora_rank, bias = False)
18 self. q_norm = nn. RMSNorm (config. q_lora_rank)
19 self. wq_b = nn. Linear (config. q_lora_rank, config. n_head * self. qk_head_dim, bias = False)
20
21 self. wkv_a = nn. Linear (config. n_embd, config. kv_lora_rank + config. qk_rope_head_dim, bias = False)
22 self. kv_norm = nn. RMSNorm (config. kv_lora_rank)
23 self. wkv_b = nn. Linear (config. kv_lora_rank, config. n_head * (config. qk_nope_head_dim + config. v_head_dim), bias = False)
24 self. kv_emb = nn. Embedding (config. vocab_size, config. kv_emb_dim)
25 self. kv_up = nn. Linear (config. kv_emb_dim, config. n_head * (config. qk_nope_head_dim + config. v_head_dim), bias = False)
26 self. ln_emb = nn. LayerNorm (config. n_head * (config. qk_nope_head_dim + config. v_head_dim))
27 self. c_proj = nn. Linear (config. n_head * config. v_head_dim, config. n_embd, bias = False)
28
29 def forward (self, x, idx, freqs_cis):
30 B, T, _ = x. size ()
31 if self. q_lora_rank == 0:
32 q = self. wq (x)
33 else:
34 q = self. wq_b (self. q_norm (self. wq_a (x)))
35
36 q = q. view (B, T, self. n_head, self. qk_head_dim)
37 q_nope, q_pe = torch. split (q, [self. qk_nope_head_dim, self. qk_rope_head_dim], dim =-1)
38 q_pe = apply_rotary_emb (q_pe, freqs_cis)
39
40 kv = self. wkv_a (x)
41 kv, k_pe = torch. split (kv, [self. kv_lora_rank, self. qk_rope_head_dim], dim =-1)
42 k_pe = apply_rotary_emb (k_pe. unsqueeze (2), freqs_cis)
43
44 q = torch. cat ([q_nope, q_pe], dim =-1)
45 kv = self. wkv_b (self. kv_norm (kv))
46 kv = self. ln_emb (kv * self. kv_up (self. kv_emb (idx)))
47 kv = kv. view (B, T, self. n_head, self. qk_nope_head_dim + self. v_head_dim)
48 k_nope, v = torch. split (kv, [self. qk_nope_head_dim, self. v_head_dim], dim =-1)
49 k = torch. cat ([k_nope, k_pe. expand (-1, -1, self. n_head, -1)], dim =-1)
50
51 k = k. view (B, T, self. n_head, -1). transpose (1, 2)
52 q = q. view (B, T, self. n_head, -1). transpose (1, 2)
53 v = v. view (B, T, self. n_head, -1). transpose (1, 2)
54
55 y = F. scaled_dot_product_attention (q, k, v, attn_mask = None, is_causal = True)
56 y = y. transpose (1, 2). contiguous (). view (B, T, -1)
57
58 y = self. c_proj (y)
59 return y
Listing 1: Implementation of Embedding Gated MLA in pseudo PyTorch style.
## F. Comparison Between Attention Mechanisms
### Comparison of Actual KV Cache Among MHA, MLA, and EG-MLA
To ensure a fair evaluation of attention mechanism efficiency, all comparisons are conducted across models with comparable overall sizes and core architectural hyperparameters—namely, the same number of layers, model hidden dimension, and attention head configuration.
Table 10 highlights the dramatic reduction in KV cache size achieved by EG-MLA compared to standard Multi-Head Attention (MHA). While MHA stores separate key and value vectors for each attention head, resulting in 18.43K elements per token, EG-MLA reduces this number by up to 94.8%, requiring as few as 0.96K elements per token in the most compressed configuration.
Among the EG-MLA variants, EG-MLA-kv64 offers a particularly favorable trade-off between performance and memory efficiency. It reduces the KV cache per token to only 1.54K elements, which constitutes a 91.6% reduction compared to standard Multi-Head Attention (MHA), and a 59.9% reduction relative to MLA. This substantial saving is achieved without altering the number of attention heads or compromising activated model parameters.
| Model Name | Total Params | Activated Params | Embed Params | n_layer / n_emb / kv_emb_dim | q_lora_rank / kv_lora_rank | qk_head_dim / v_head_dim | KV Cache / Token (#Element) |
| --- | --- | --- | --- | --- | --- | --- | --- |
| MHA-Base | 124M | 124M | 0M | 12/768/- | - | 64 | 18.43K |
| MLA-Base | 120M | 120M | 0M | 12/768/256 | 0/256 | 64+64 / 64 | 3.84K |
| EG-MLA-Base-kv256 | 279M | 125M | 154M | 12/768/256 | 0/256 | 64+64 / 64 | 3.84K |
| EG-MLA-Base-kv128 | 276M | 122M | 154M | 12/768/256 | 0/128 | 64+64 / 64 | 2.30K |
| EG-MLA-Base-kv64 | 274M | 120M | 154M | 12/768/256 | 0/64 | 64+64 / 64 | 1.54K |
| EG-MLA-Base-kv16 | 273M | 119M | 154M | 12/768/256 | 0/16 | 64+64 / 64 | 0.96K |
Table 10: Comparison of actual KV cache usage among Multi-Head Attention (MHA), Multi-Head Latent Attention (MLA), and our proposed Embedding Gated MLA (EG-MLA). EG-MLA achieves comparable or stronger performance while significantly reducing the KV cache per token, thanks to its compressed latent representation and token-specific gating.
## G. More Scaling Experiment
Figure 7 shows the difference in validation loss between the EG-MLA-1.2B model and the baseline MLA-1.2B model over the course of training. The x-axis represents the number of tokens processed (in billions), while the y-axis indicates the validation loss difference. A negative value implies that EG-MLA-1.2B performs better (i.e., lower validation loss) than the baseline at that point in training. The trend highlights how the enhancements in EG-MLA-1.2B translate into improved generalization throughout the training trajectory.
<details>
<summary>figures/EG-MLA-1.2B_vs_MLA-1.2B_val_loss_diff.png Details</summary>
### Visual Description
## Line Chart: Validation Loss Difference vs Tokens Trained
### Overview
The chart illustrates the relationship between the number of tokens trained (in billions) and the validation loss difference over time. The y-axis represents the validation loss difference (negative values indicate improvement), while the x-axis tracks training progress in billions of tokens. The data shows a general trend of improvement with increasing training data, though with notable fluctuations.
### Components/Axes
- **X-axis (Horizontal)**: "Tokens Trained (B)" with markers at 0, 10, 20, 30, 40, and 50 billion tokens.
- **Y-axis (Vertical)**: "Validation Loss Difference" ranging from -0.046 to -0.032 in increments of 0.002.
- **Legend**: Located in the top-right corner, indicating a single data series represented by blue circular markers connected by a line.
- **Data Points**: Blue dots with error bars (not explicitly labeled but visually present).
### Detailed Analysis
- **Initial Phase (0–10B tokens)**:
- Validation loss difference starts at approximately **-0.046** at 0 tokens.
- Rapid improvement occurs, reaching **-0.042** by 5B tokens.
- Slight plateau observed between 5–7B tokens (~ -0.043 to -0.042).
- **Mid-Phase (10–30B tokens)**:
- Continued improvement to **-0.038** by 15B tokens.
- Notable volatility between 15–20B tokens, with a local minimum of **-0.039** at 18B tokens.
- Sharp improvement to **-0.036** by 25B tokens, followed by stabilization around **-0.035** to **-0.037**.
- **Late Phase (30–50B tokens)**:
- Gradual improvement to **-0.034** by 40B tokens.
- Increased volatility observed after 40B tokens, with a peak of **-0.032** at 45B tokens.
- Final value at 50B tokens: **-0.035**.
### Key Observations
1. **Overall Trend**: Validation loss difference improves (becomes less negative) as tokens trained increase, suggesting better model generalization.
2. **Volatility**: Significant fluctuations occur in the 15–20B and 40–50B token ranges, potentially indicating overfitting, data complexity, or training instability.
3. **Inflection Points**:
- Steepest improvement between 5–10B tokens.
- Local minimum at 18B tokens (-0.039).
- Peak improvement at 45B tokens (-0.032).
### Interpretation
The data demonstrates that increasing training data generally enhances model performance, as evidenced by the negative validation loss difference trending toward zero. However, the volatility in later stages suggests potential challenges:
- **Overfitting**: Sharp improvements followed by reversals may indicate the model is memorizing noise in the training data.
- **Data Complexity**: Fluctuations in the 40–50B range could reflect difficulties in generalizing to edge cases or domain shifts.
- **Training Dynamics**: The mid-phase plateau (15–20B tokens) might correspond to architectural bottlenecks or optimization challenges.
The chart underscores the importance of balancing dataset size with regularization techniques to maintain stable generalization. The final validation loss difference of -0.035 at 50B tokens suggests moderate improvement but leaves room for optimization through hyperparameter tuning or architectural adjustments.
</details>
Figure 7: Difference in validation loss between EG-MLA-1.2B and MLA-1.2B models as a function of tokens trained.
| Model Name | Model Size | Embed Size | piqa | arc-c | arc-e | hella | winog | siqa | mmlu | obqa | boolq | race | truth | lmb.o | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| EG-MLA-Large-ctx2048-kv128emb512 | 951M | 309M | 75.35 | 37.20 | 70.33 | 42.50 | 53.83 | 90.80 | 27.22 | 26.80 | 57.55 | 33.01 | 29.61 | 39.72 | 48.66 |
| MLA-Large-ctx2048-kv512 | 645M | 0M | 74.54 | 37.88 | 69.70 | 42.16 | 55.01 | 90.40 | 27.34 | 27.20 | 62.97 | 34.83 | 31.12 | 41.24 | 49.53 |
| EGA-XLarge-ctx4096-kv128emb512 | 1800M | 618M | 75.41 | 38.74 | 72.85 | 44.88 | 56.83 | 91.00 | 26.92 | 27.80 | 59.45 | 33.97 | 31.27 | 40.25 | 49.95 |
| MLA-XLarge-ctx4096-kv512 | 1188M | 0M | 75.52 | 38.82 | 73.02 | 44.15 | 56.83 | 92.20 | 28.42 | 27.20 | 53.15 | 34.26 | 33.76 | 38.02 | 49.61 |
| EG-MLA-Large-ctx2048-kv256emb512 | 961M | 309M | 74.92 | 37.12 | 71.55 | 42.79 | 54.38 | 90.70 | 25.93 | 23.80 | 58.96 | 33.21 | 29.24 | 41.98 | 48.71 |
Table 11: Accuracy results of more models across 12 benchmark tasks. kv refers to KV cache dimension and emb refers to embedding dimension.All models were trained for a total of 50 billion tokens.
| Model Name | Model Size | Embed Size | ClimbMix-val | wiki | lambada | Avg. |
| --- | --- | --- | --- | --- | --- | --- |
| EG-MLA-Large-ctx2048-kv128emb512 | 951M | 309M | 10.79 | 6.17 | 19.57 | 12.18 |
| MLA-Large-ctx2048-kv512 | 645M | 0M | 11.02 | 6.31 | 18.95 | 12.09 |
| EGA-XLarge-ctx4096-kv128emb512 | 1800M | 618M | 10.35 | 2.94 | 19.01 | 10.77 |
| MLA-XLarge-ctx4096-kv512 | 1188M | 0M | 10.71 | 3.00 | 22.99 | 12.23 |
| EG-MLA-Large-ctx2048-kv256emb512 | 961M | 309M | 10.58 | 6.14 | 17.95 | 11.56 |
Table 12: Perplexity results of more models across 12 benchmark tasks. kv refers to KV cache dimension and emb refers to embedding dimension.All models were trained for a total of 50 billion tokens.
## H. Compression through Layer Grouping
To push the optimization frontier and further curtail computational and memory costs, we explore the reuse of both the KV-projection parameters and their intermediate activations across transformer layers. Specifically, we introduce a hyper-parameter $layergroupsize(lgz)$ that dictates how many consecutive transformer layers share a single parameter set rather than maintaining independent parameters for each layer. During the forward pass, all layers within the same group directly reuse the shared parameters and cached intermediate results, eliminating redundant computation and improving training and inference efficiency.
Empirical results show that when $lgz$ = 2, the performance of model remains comparable to the baseline. However, once $lgz$ exceeds 4, a significant performance degradation is observed. Therefore, $lgz$ = 4 represents the compression limit under current settings, and we recommend $lgz$ = 3 as a balanced trade-off between efficiency gains and accuracy loss.
| Model Name | piqa | arc-c | arc-e | hella | winog | siqa | mmlu | obqa | boolq | race | truth | lmb.o | Avg. | Diff (%) |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| MLA-Large-ctx2048-kv512 | 74.54 | 37.88 | 69.70 | 42.16 | 55.01 | 90.40 | 27.34 | 27.20 | 62.97 | 34.83 | 31.12 | 41.24 | 49.53 | — |
| EG-MLA-Large-ctx2048-kv256emb512 | 74.92 | 37.12 | 71.55 | 42.79 | 54.38 | 90.70 | 25.93 | 23.80 | 58.96 | 33.21 | 29.24 | 41.98 | 48.71 | — |
| EG-MLA-Large-ctx2048-kv256emb512lgz2 | 74.76 | 37.03 | 69.57 | 42.56 | 52.33 | 90.50 | 25.65 | 27.80 | 59.79 | 31.87 | 30.59 | 39.90 | 48.53 | -0.37 |
| EG-MLA-Large-ctx2048-kv256emb512lgz4 | 75.03 | 36.18 | 69.32 | 41.61 | 52.25 | 91.60 | 27.70 | 26.40 | 55.69 | 32.25 | 29.27 | 40.50 | 48.15 | -1.15 |
| EG-MLA-Large-ctx2048-kv256emb512lgz6 | 74.48 | 36.35 | 70.58 | 40.69 | 52.80 | 90.10 | 26.98 | 24.60 | 58.10 | 31.20 | 29.20 | 38.48 | 47.80 | -1.87 |
Table 13: Accuracy results across 12 benchmark tasks under different layer group sizes (lgz). Diff is the relative change (%) compared to EG-MLA-Large-ctx2048-kv256emb512.
| Model Name | ClimbMix-val | wiki | lambada | Avg. | Diff (%) |
| --- | --- | --- | --- | --- | --- |
| MLA-Large-ctx2048-kv512 | 11.02 | 6.31 | 18.95 | 12.09 | — |
| EG-MLA-Large-ctx2048-kv256emb512 | 10.58 | 6.14 | 17.95 | 11.56 | — |
| EG-MLA-Large-ctx2048-kv256emb512lgz2 | 10.78 | 6.23 | 18.69 | 11.90 | -2.94 |
| EG-MLA-Large-ctx2048-kv256emb512lgz4 | 11.05 | 6.38 | 20.19 | 12.54 | -8.47 |
| EG-MLA-Large-ctx2048-kv256emb512lgz6 | 11.27 | 6.51 | 22.77 | 13.52 | -16.98 |
Table 14: Perplexity results across 12 benchmark tasks under different layer group sizes (lgz). Diff is the relative change (%) compared to EG-MLA-Large-ctx2048-kv256emb512.