在大語(yǔ)言模型蓬勃發(fā)展的背景下，Transformer 架構(gòu)依然是不可替代的核心組件。盡管其自注意力機(jī)制存在計(jì)算復(fù)雜度為二次方的問(wèn)題，成為眾多研究試圖突破的重點(diǎn)，但 Transformer 在推理時(shí)靈活建模長(zhǎng)距離上下文的能力，使得許多線性復(fù)雜度的替代方案（如 RNN、Linear Attention、SSM 等）難以真正取代它的地位。

尤其是在大語(yǔ)言模型廣泛采用 decoder-only 架構(gòu)之后，自注意力機(jī)制的重要性進(jìn)一步凸顯。然而，這種機(jī)制也帶來(lái)新的挑戰(zhàn)：推理過(guò)程中每一步都需要訪問(wèn) Key-Value（KV）緩存，該緩存的大小隨著生成序列長(zhǎng)度線性增長(zhǎng)，逐漸成為影響推理效率的關(guān)鍵瓶頸。隨著模型參數(shù)維度不斷擴(kuò)大，KV 緩存所需的顯存和帶寬開(kāi)銷顯著上升，限制了模型的推理長(zhǎng)度與可支持的 batch size。

值得一提的是，近期由 DeepSeek 團(tuán)隊(duì)提出的 MLA 機(jī)制，通過(guò)在隱空間維度對(duì) KV 緩存進(jìn)行壓縮，顯著提升了推理效率，推動(dòng)了大模型在低資源場(chǎng)景下的高效部署。但隨著生成序列的持續(xù)增長(zhǎng)，時(shí)間維度的冗余信息也逐漸暴露，壓縮其所帶來(lái)的潛力亟待挖掘。然而，如何在保持性能的前提下壓縮時(shí)間維度，一直受到增量式推理復(fù)雜性的限制。

為此，劍橋大學(xué)機(jī)器智能實(shí)驗(yàn)室最新提出了 Multi-head Temporal Latent Attention（MTLA），首次將時(shí)序壓縮與隱空間壓縮相結(jié)合，在 KV 緩存的兩個(gè)維度上同時(shí)施加時(shí)空壓縮策略。MTLA 利用超網(wǎng)絡(luò)動(dòng)態(tài)融合相鄰時(shí)間步的信息，并設(shè)計(jì)了步幅感知的因果掩碼以確保訓(xùn)練與推理的一致性，在顯著降低推理顯存與計(jì)算成本的同時(shí)，保持甚至略優(yōu)于傳統(tǒng)注意力機(jī)制的模型性能，為大語(yǔ)言模型推理效率的提升提供了新的解決思路。

• 論文標(biāo)題：Multi-head Temporal Latent Attention
• 論文地址：https://arxiv.org/pdf/2505.13544
• 項(xiàng)目地址：https://github.com/D-Keqi/mtla

現(xiàn)有方法的局限與 MTLA 的突破

在構(gòu)建大語(yǔ)言模型時(shí)，KV 緩存帶來(lái)的顯存與計(jì)算開(kāi)銷問(wèn)題早已受到廣泛關(guān)注。當(dāng)前主流的大模型通常采用基于自注意力的 Grouped-Query Attention（GQA）機(jī)制，對(duì)標(biāo)準(zhǔn) Transformer 中的 Multi-Head Attention（MHA）進(jìn)行改進(jìn)。GQA 通過(guò)減少 Key/Value 頭的數(shù)量來(lái)減小 KV 緩存的規(guī)模，具體做法是將多個(gè) Query 頭分組，每組共享同一個(gè) KV 頭。

當(dāng) GQA 的組數(shù)等于 Query 頭數(shù)量時(shí)，其等價(jià)于標(biāo)準(zhǔn) MHA；而當(dāng)組數(shù)為 1 時(shí)，即所有 Query 頭共享同一組 KV，這種極端形式被稱為 Multi-Query Attention（MQA）。雖然 MQA 極大地減少了顯存占用，但顯著影響模型性能；相比之下，GQA 在效率與效果之間取得了更好的平衡，因此成為當(dāng)前大語(yǔ)言模型中最常見(jiàn)的注意力變體。

與此不同，DeepSeek 團(tuán)隊(duì)提出的 Multi-head Latent Attention（MLA）采用了另一種思路：不減少頭的數(shù)量，而是在隱空間中壓縮 KV 的特征維度。實(shí)驗(yàn)結(jié)果表明，MLA 相較于 GQA 表現(xiàn)出更優(yōu)的性能與效率。然而，這種壓縮方式仍存在上限，為了維持模型性能，隱空間維度的壓縮幅度不能過(guò)大，因此 KV 緩存的存儲(chǔ)開(kāi)銷依然是限制模型推理效率的一大瓶頸。

除了在隱空間對(duì) KV 緩存進(jìn)行壓縮之外，時(shí)間維度也是一個(gè)極具潛力但尚未充分挖掘的方向。隨著生成序列變得越來(lái)越長(zhǎng)，KV 緩存中在時(shí)間軸上的信息冗余也日益明顯。然而，由于自注意力機(jī)制在生成時(shí)通常采用自回歸的增量推理模式，KV 緩存與每一個(gè)生成的 token 是一一對(duì)應(yīng)的，這使得在保持模型性能的前提下壓縮時(shí)間維度成為一項(xiàng)挑戰(zhàn)，也導(dǎo)致了該方向長(zhǎng)期缺乏有效解決方案。

MTLA 的提出正是對(duì)這一空白的回應(yīng)。它通過(guò)引入時(shí)間壓縮機(jī)制和步幅感知的因果掩碼，巧妙解決了訓(xùn)練與推理行為不一致的問(wèn)題，在保持高效并行訓(xùn)練能力的同時(shí)，實(shí)現(xiàn)了推理過(guò)程中的 KV 時(shí)間壓縮。進(jìn)一步地，MTLA 還結(jié)合了 MLA 的隱空間壓縮策略，從空間與時(shí)間兩個(gè)維度同時(shí)優(yōu)化 KV 緩存的表示，將自注意力機(jī)制的效率推向了新的高度。

MTLA 的核心技術(shù)與訓(xùn)練策略

在增量推理階段，MTLA 會(huì)對(duì)經(jīng)過(guò)隱空間壓縮后的 KV 緩存進(jìn)行時(shí)間維度的增量式合并，進(jìn)一步壓縮存儲(chǔ)空間。上圖展示了該過(guò)程的示意，并與標(biāo)準(zhǔn)的 MHA 進(jìn)行了對(duì)比。

以時(shí)間壓縮率 s=2 為例，每?jī)蓚€(gè)相鄰的 KV 緩存將合并為一個(gè)。在生成第一個(gè)字符時(shí)，KV 緩存長(zhǎng)度為 1；生成第二個(gè)字符后，新生成的 KV 與前一個(gè)被合并，KV 緩存長(zhǎng)度仍然保持為 1。這種動(dòng)態(tài)合并機(jī)制有效壓縮了時(shí)間維度上的冗余信息。

然而，這也帶來(lái)了并行訓(xùn)練上的挑戰(zhàn)：雖然兩個(gè)時(shí)間步的 KV 緩存長(zhǎng)度相同，但它們所包含的信息不同，若不加以區(qū)分，容易導(dǎo)致訓(xùn)練與推理行為不一致。

MTLA 通過(guò)一種優(yōu)雅的方式解決了這一問(wèn)題。正如下圖所示，在訓(xùn)練階段，MTLA 保留了所有中間狀態(tài)的 KV 表達(dá)，并引入了步幅感知因果掩碼（stride-aware causal mask），確保每個(gè) query 在訓(xùn)練時(shí)訪問(wèn)到與推理階段一致的 KV 區(qū)域，從而準(zhǔn)確模擬增量推理中的注意力行為。

得益于這一設(shè)計(jì)，MTLA 能夠像標(biāo)準(zhǔn)注意力機(jī)制一樣通過(guò)矩陣乘法實(shí)現(xiàn)高效并行計(jì)算，在保持訓(xùn)練效率的同時(shí)完成對(duì)時(shí)間維度的壓縮。

此外，MTLA 還引入了解耦的旋轉(zhuǎn)位置編碼（decoupled RoPE）來(lái)建模位置信息，并對(duì)其進(jìn)行了時(shí)間維度上的壓縮，進(jìn)一步提升了整體效率。

值得強(qiáng)調(diào)的是，MTLA 不僅是一種更高效的自注意力機(jī)制，它還具備極強(qiáng)的靈活性與可調(diào)性。例如，當(dāng)將時(shí)間壓縮率 s 設(shè)置得足夠大時(shí)，MTLA 在推理過(guò)程中幾乎只保留一個(gè) KV 緩存，這種形式本質(zhì)上就退化為一種線性序列建模方法。換句話說(shuō)，線性序列建模可以被視為 MTLA 的極端情況，MTLA 在注意力機(jī)制與線性模型之間架起了一座橋梁。

然而，在許多復(fù)雜任務(wù)中，傳統(tǒng)注意力機(jī)制所具備的二次計(jì)算復(fù)雜度雖然代價(jià)高昂，卻提供了更強(qiáng)的建模能力。因此，MTLA 所引入的 “可調(diào)時(shí)間壓縮率 s” 這一設(shè)計(jì)思路，恰恰為模型提供了一個(gè)在效率與性能之間靈活權(quán)衡的可能空間。

MTLA 的卓越性能

MTLA 在一系列任務(wù)中展現(xiàn)了出色的性能，包括語(yǔ)音翻譯，文本摘要生成，語(yǔ)音識(shí)別和口語(yǔ)理解。例如在語(yǔ)音翻譯中，MTLA 在保持與標(biāo)準(zhǔn) MHA 相當(dāng)?shù)姆g質(zhì)量的同時(shí)，實(shí)現(xiàn)了超過(guò) 5 倍的推理速度提升，并將推理過(guò)程中的 GPU 顯存占用降低了超過(guò) 8 倍。

值得注意的是，僅當(dāng)時(shí)間壓縮率 s=2 時(shí)，MTLA 對(duì) KV 緩存的壓縮程度就已經(jīng)與 MQA 相當(dāng)，且在模型性能上更具優(yōu)勢(shì)。而相比之下，MQA 所采用的減少 KV 頭數(shù)量的方法已達(dá)上限，而 MTLA 還有進(jìn)一步的空間。

未來(lái)發(fā)展

MTLA 具備在大規(guī)模場(chǎng)景中部署的顯著潛力，尤其是在大語(yǔ)言模型參數(shù)規(guī)模不斷擴(kuò)大、以及思維鏈等技術(shù)推動(dòng)下生成序列日益增長(zhǎng)的背景下，對(duì) KV 緩存進(jìn)行時(shí)空壓縮正是緩解推理開(kāi)銷的關(guān)鍵手段。在這樣的趨勢(shì)下，MTLA 有望成為未來(lái)大語(yǔ)言模型中自注意力模塊的重要替代方案。

當(dāng)然，與 DeepSeek 提出的 MLA 類似，MTLA 相較于 GQA 和 MQA，在工程落地方面的改動(dòng)不再是簡(jiǎn)單的一兩行代碼可以實(shí)現(xiàn)的優(yōu)化。這也意味著要將其大規(guī)模應(yīng)用到現(xiàn)有 LLM 框架中，還需要來(lái)自社區(qū)的持續(xù)推動(dòng)與協(xié)同開(kāi)發(fā)。

為促進(jìn)這一過(guò)程，MTLA 的實(shí)現(xiàn)代碼已全面開(kāi)源，希望能夠?yàn)檠芯空吲c工程實(shí)踐者提供便利，共同推動(dòng)高效注意力機(jī)制在大模型時(shí)代的落地與普及。