Multi-head Temporal Latent Attention

TL;DR

MTLA extends DeepSeek’s Multi-head Latent Attention (MLA) with a second axis of KV-cache compression: the temporal dimension. A hyper-network dynamically merges every $s$ temporally adjacent low-rank latent KV vectors into one, so the cache shrinks by a further factor $s$ on top of MLA’s low-rank compression. A stride-aware causal mask makes parallel training match the incremental-decoding attention pattern, and decoupled RoPE keys are compressed along the same temporal axis. On English-German speech translation MTLA at $s=2$ matches or slightly beats MHA in BLEU while delivering 4.29× inference speedup and 6.58× lower GPU memory; at $s=4$ the speedup is 5.78× at 9.71× lower memory with a small quality drop. NeurIPS 2025; arrived via the open-source D-Keqi/mtla toolkit that also ships a customized FlashAttention-2 backend for MTLA inference.

Key claims

The KV cache grows linearly with sequence length and dominates decode-time memory bandwidth in long-context auto-regressive Transformers; existing MHA / MQA / GQA / MLA variants compress the head and latent dimensions but leave the temporal dimension uncompressed, which is the gap MTLA targets [§1, §2].
MTLA compresses the temporal dimension by dynamically merging every $s$ adjacent low-rank latent vectors $\mathbf c_t^{KV}$ into one, using merging weights $\mathbf W_t$ produced by a hyper-network conditioned on the latent input plus positional embedding (so the merging adapts per sample rather than relying on static parameters) [§4, §4.1, Eq. 13].
The KV-cache-length / processed-sequence-length mismatch breaks naive parallel training; the proposed stride-aware causal mask zeros out attention only when $j \le \lceil i/s \rceil$ or $j = \lceil i/s \rceil$ (the “currently being filled” cache slot is attendable for its own block), matching the incremental-decoding pattern exactly [§4.2, Fig. 2(c)].
Decoupled RoPE keys (used by MLA to keep the rotated dimension out of the absorbed projection) can be compressed along the same temporal axis with the same stride-aware mask; this lets the original-length $\mathbf k^R$ be used during training in place of the compressed $\mathbf c^R$ , simplifying the training pipeline [§4.3, Eq. 17].
With default hyper-parameters ( $d_c = 4 d_h$ , decoupled-RoPE head dim $d_h^R = d_h/2$ , $s=2$ ), MTLA stores on average $\frac{9}{4s} = \frac{9}{8}$ KV-cache elements per token per layer — close to MQA’s 1 — but with quality matching MHA [§4.3].
On MuST-C En–De speech translation, MTLA at $s=2$ scores 23.28 BLEU vs MHA 23.18, MLA 22.97, MQA 22.70, GQA 22.75 — competitive with MHA while compressing the cache to MQA-equivalent footprint [Table 1, Table 5].
Compared to MHA the same MTLA- $s{=}2$ achieves 4.29× wall-clock inference speedup on RTX 6000 Ada and 6.58× lower average GPU memory; pushing $s$ to 3 and 4 trades small BLEU for further gains, reaching 5.78× speedup / 9.71× memory reduction at $s=4$ [Table 1].
The pattern transfers across tasks: XSum text summarization (ROUGE-L 23.60 vs MHA 23.33, 3.35× speedup, 7.34× memory), AMI ASR (WER 12.66 vs MHA 12.98, 3.75× speedup, 7.41× memory), and SLURP intent classification (86.80% vs MHA 86.83%, 2.53× speedup, 7.01× memory) [Tables 2, 3, 4].
Per-token decode complexity drops from $O(N)$ to $O(N/s)$ , but because the feed-forward modules are unchanged, doubling $s$ does not halve wall-clock; observed gains are 1.48× at $s=2$ and 1.79× memory reduction over MLA [§6.1].
The “compressing redundant historical KV information may sometimes benefit performance” observation — MTLA- $s{=}2$ beating MHA on ST BLEU and on long-context-friendly tasks — is presented as an emergent regularization side-effect of the temporal compression [§6.1].

Method

MTLA inherits MLA’s low-rank compression — input $\mathbf h_t$ projects to a shared latent $\mathbf c_t^{KV} \in \mathbb R^{d_c}$ with $d_c \ll H d_h$ , which is what gets cached. On top of MLA, MTLA introduces three pieces. (1) A hyper-network takes $\mathbf c_t^{KV}$ (element-wise multiplied with positional embedding) through a linear layer to produce merging weights $\mathbf W_t \in \mathbb R^s$ , which combine $s$ temporally adjacent latents into one merged latent stored in the cache. During inference, only every $s$ -th step writes a new cache slot; intermediate steps update the most recent (still-being-filled) slot in place. (2) A stride-aware causal mask for training: instead of pre-downsampling the latent sequence (which would not match inference behavior because incremental decoding lets a query at step $t$ attend to its own block’s partially-filled latent), the mask is set to 0 only when $j \le \lceil i/s \rceil$ or $j = \lceil i/s \rceil$ , replicating the incremental pattern under matrix-multiply parallel training. (3) Decoupled-RoPE keys $\mathbf k_t^R$ are compressed along the same temporal axis with the same mask, so the same training-time substitution works for both the absorbed and the rotated halves of the attention computation. Equations 11–17 in the paper specify the full forward pass; matrix-multiply absorption ( $\mathbf W_q^{UK}$ into $\mathbf W_q$ , $\mathbf W^{UV}$ into the output projection) is preserved from MLA so the explicit per-step $\mathbf k, \mathbf v$ never need to be reconstructed.

The companion GitHub repo D-Keqi/mtla ships an MTLA module, recipes for MuST-C ST / XSum summarization / AMI ASR / SLURP SLU built on Fairseq, HuggingFace Transformers integration (LlamaMTLAConfig, LlamaMTLAForCausalLM), and a forked FlashAttention-2 kernel that handles the stride-aware mask at inference; fp16/bf16 only.

Results

Speech translation (MuST-C En–De): MHA 23.18 BLEU / 281.3s / 18646 MiB; MLA 22.97 / 97.0s / 5065 MiB; MTLA- $s{=}2$ 23.28 / 65.6s / 2835 MiB (4.29× speedup, 6.58× memory). MTLA- $s{=}3$ : 23.25 / 52.7s / 2251 MiB. MTLA- $s{=}4$ : 23.05 / 48.7s / 1921 MiB [Table 1].
Text summarization (XSum): MTLA- $s{=}2$ ROUGE-1/2/L 29.14/9.79/23.60 vs MHA 28.83/9.67/23.33; 3.35× speedup, 7.34× memory reduction [Table 2].
Speech recognition (AMI): MTLA- $s{=}2$ WER 12.66 vs MHA 12.98, MLA 12.67; 3.75× speedup, 7.41× memory reduction [Table 3].
Spoken language understanding (SLURP intent): MTLA- $s{=}2$ 86.80% accuracy vs MHA 86.83%; 2.53× speedup, 7.01× memory reduction [Table 4].
vs MQA/GQA at matched cache footprint (Table 5): MTLA- $s{=}2$ beats MQA (22.70 BLEU) and GQA (22.75) by ~0.5 BLEU on ST at roughly equivalent cache size, with faster inference than both because the low-rank latent skips explicit per-head $\mathbf k,\mathbf v$ reconstruction. At $s{=}4$ MTLA significantly outperforms MQA on BLEU ( $p < 0.05$ , SacreBLEU) while compressing further than MQA can ([§6.3]).
Experiments are decoder-only (encoder output prepended to the self-attention input, no cross-attention), trained from scratch at 512 / 8-head decoder, $d_c = 256$ , decoupled-RoPE dim 32, on the four task datasets — i.e. no LLM-scale pretraining was tested. The recipe is reproducible from the open-source toolkit.

Why it’s interesting

MTLA is the first filed paper to compress the temporal dimension of the KV cache — the third axis after MQA/GQA’s head-count compression and MLA’s per-token latent compression. It’s a clean architectural variant that the wiki’s LLM Inference Efficiency page has been collecting indirect attacks on (speculative decoding, DMS eviction, weight quantization, prefix-cacheable diffusion LMs) without a frontal architectural play — MTLA is the architectural lever. Sits naturally next to Inference-Time Hyper-Scaling with KV Cache Compression, which compresses an existing cache by eviction with delayed-eviction sliding window; MTLA instead never produces the redundant cache entries in the first place via learned temporal merging, and provides a complementary recipe — eviction-style methods retrofit; MTLA requires training the attention from scratch. The How Attention Got So Efficient [GQA / MLA / DSA] explainer covers the MHA → GQA → MLA → DSA progression that MTLA extends along a previously-unused axis. Caveat: experiments are at small (decoder-512) scale on speech/summarization tasks rather than LLM-pretraining scale, so the “no quality loss” claim still needs an LLM-scale replication — the released HuggingFace integration makes that test cheap to run.