Multi-matrix Factorization Attention (2025)

Abstract: We propose novel attention architectures, Multi-matrix Factorization Attention (MFA) and MFA-Key-Reuse (MFA-KR). Existing variants for standard Multi-Head Attention (MHA), including SOTA methods like MLA, fail to maintain as strong performance under stringent Key-Value cache (KV cache) constraints. MFA enhances model capacity by efficiently scaling up both the number and dimension of attention heads through low-rank matrix factorization in the Query-Key (QK) circuit. Extending MFA, MFA-KR further reduces memory requirements by repurposing the key cache as value through value projection re-parameterization. MFA’s design enables strong model capacity when working under tight KV cache budget, while MFA-KR is suitable for even harsher KV cache limits with minor performance trade-off. Notably, in our extensive and large-scale experiments, the proposed architecture outperforms MLA and performs comparably to MHA, while reducing KV cache usage by up to 56% and 93.7%, respectively.

## introduction the decoder-only transformer with standard multi-head attention (mha) @xcite has become the de facto architecture for large language models. its autoregressive nature enables the reuse of cached attention key-value tensors (kv cache) from previous tokens, significantly relieving the computation overhead during the stepby-step decoding @xcite . however, the kv cache memory footprint scales linearly with both batch size and sequence length, leading to large amount of memory occupancy and traffic, which becomes the primary bottleneck during the decoding phase of llm @xcite . to address these challenges, multi-query attention (mqa) and grouped query attention (gqa) reduce kv cache usage by sharing key and value projections across heads @xcite . similarly, multi-head latent attention (mla) applies low-rank compression to key and value projections and only caches the latents @xcite . however, all methods fail to match mha's performance under stringent kv cache budgets @xcite , as the added constraints on key and value projections limit the capacity of the attention module. driven by these limitations, we analyze the modeling capacity in attention mechanisms and present a unified perspective on existing mha variants. our analysis reveals that the number and dimension of attention heads are critical for maintaining modeling capacity-an under-explored design aspect in current methods @xcite @xcite . this insight highlights the need to scale these factors efficiently to mitigate the capacity degradation caused by existing kv cache-saving techniques, pushing attention modules closer to their theoretical upper bound. inspired by this understanding, we propose multi-matrix factorization attention (mfa), a novel attention module, along with its variant, mfa-key-reuse (mfa-kr). these attention modules are specifically designed to enhance modeling capacity under strict kv cache constraints. specifically, mfa employs a low-rank matrix factorization in the query-key (qk) circuit @xcite , enabling parameter-efficient scaling of both the number and dimension of heads without excessive kv cache usage. building on mfa, mfa-kr reuses the key cache as value through a re-parameterized value projection with original key projection and a light weight gated projection. this minor modification cuts kv cache usage by an additional 50% with negligible performance trade-offs. moreover, methods like mla add complexity to support widely-adopted position embedding (i.e. rope), while our proposed mfa family naturally fit in current llm training and inference ecosystems, ensuring practical adoption without introducing additional architectural complexity. we conduct extensive experiments to evaluate the performance and kv cache efficiency of mfa and mfa-kr, alongside detailed ablation studies on their design. impressively, our proposed attention architecture is the only approach that performs comparably to standard mha in terms of accuracy while adhering to strict kv cache constraints. specifically, in a 7b parameter model trained on 1t tokens, mfa and mfa-kr reduce kv cache usage by up to 93.7% while achieving superior or comparable benchmark accuracies compared to mha. ## background: capacity analysis of attention in order to delimit the scope of our analysis, we introduce the concept of generalized multi-head attention (gmha). it encompasses all multi-head mechanisms with linear query-key (qk) and valueoutput (vo) circuits, and per-head softmax attention. the qk circuit determines how information propagates between entities, and the vo circuit dictates how information is transformed @xcite . fundamentally, gmha can be described and analyzed using inference formulation and factorization formulation. the inference formulation highlights how keys and values are computed and cached during inference, and factorization formulation clarifies the model's capacity by interpreting qk and vo matrices as low-rank factorizations. this offers a unified perspective on how different factorization strategies mediate the trade-off between model capacity and efficiency. within this framework, we identify fully parameterized bilinear attention (fpba) as the upper bound of capacity. mha and its variants can be regarded as a low-rank decomposition of fpba, making fpba a unified theoretical reference point for analysis. building on this understanding, we propose general design principles for constructing efficient and effective attention modules. these principles inform the design of the multi-matrix factorization attention (mfa) mechanism, which is introduced in the next section. ## multi-matrix factorization attention building upon the analysis in the previous section, we arrive at the general design objective for efficient and effective attention module: to find a matrix factorization scheme that minimizes parameter and kv cache size while pushing the model's capacity as close as possible to that of fpba. following these principles, we introduce multimatrix factorization attention (mfa), incorporating three key design strategies: (1) increasing the number and dimension of heads to minimize the amount of channel sharing in the propagation process and to provide greater expressive freedom for each head; (2) applying aggressive low-rank matrix factorizations on w n to enhance parameter efficiency as the model scales; (3) utilizing singlekey-and-value-head techniques to maintain minimal kv cache usage. the inference and factorization expressions of mfa are given by: ## experiments we evaluate mfa for large language models from the following perspectives. first, we compare mfa and mfa-kr to mha at 7b-scale moe models with 1t training tokens on benchmark accuracies and kv cache usage. second, we present the loss and kv cache curves of mfa and mfa-kr on increasing training scales. third, we conduct extensive comparison with existing architecture variants and demonstrate the advantage of mfa and mfa-kr. finally, we present studies on various design choices and validate the compatibility with different position embeddings ## related works notable efforts have focused on architectural modifications to minimize kv cache usage besides mqa, gqa and mla which we elaborated in previous section. cla @xcite and @xcite length, such as ssms @xcite or linear attention @xcite . this reduces the cache state size significantly in extremely long-context region, and can be combined with our proposed mfa and mfa-kr in hybrid manner. another active area of research seeks to boost the capacity of attention modules. @xcite indentifies the dimension of each head in mha bottlenecks the capacity of attention module, and the situation may become worse if adhere to current high weight decay training recipe @xcite . other works like talking-head attention @xcite and dcmha @xcite try to enable information exchange between heads to augment model capacity. though potential performance gain can be achieved, this modification is not compatible with commonly used flash attention @xcite , limiting the scaling up of these architectures. parameter efficiency in transformer models has been extensively studied, especially in finetuning domain @xcite . there are also works focusing on pretraining parameter efficiency like lpa @xcite ; however, they do not investigate the effects under limited kv cache budget. ## conclusions we present multi-matrix factorization attention (mfa) and its variant mfa-key-reuse (mfa-kr) as scalable solutions to achieve superior perfor-mance while drastically reducing kv cache requirements. our experiments demonstrate that mfa achieves superior benchmark accuracies with up to 87.5% less kv cache compared to mha, while mfa-kr pushes memory efficiency further by halving kv cache requirements with minimal trade-offs.