Multi-head Temporal Latent Attention

TL;DR

MTLA advances decoder-only Transformer efficiency by compressing the KV cache along the temporal dimension with a learnable hyper-network that merges adjacent latent vectors. It extends MLA with temporal compression and introduces a stride-aware causal mask to align training with incremental decoding, while integrating decoupled RoPE. Across speech translation, speech recognition, speech understanding, and text summarisation, MTLA achieves competitive accuracy with substantial inference speedups and GPU memory reductions compared to MHA and MLA, including strong gains even when paired with FlashAttention-2. The approach shows promise for scalable long-context decoding and potential applicability to large language models, offering a practical path toward more memory-efficient, faster inference without sacrificing performance.

Abstract

While Transformer self-attention offers strong parallelism, the Key-Value (KV) cache grows linearly with sequence length and becomes a bottleneck for inference efficiency. Multi-head latent attention was recently developed to compress the KV cache into a low-rank latent space. This paper proposes Multi-head Temporal Latent Attention (MTLA), which further reduces the KV cache size along the temporal dimension, greatly lowering the memory footprint of self-attention inference. MTLA employs a hyper-network to dynamically merge temporally adjacent KV cache vectors. To address the mismatch between the compressed KV cache and processed sequence lengths, a stride-aware causal mask is proposed to ensure efficient parallel training and consistency with inference behaviour. Experiments across tasks, including speech translation, speech recognition, speech understanding and text summarisation, demonstrate that MTLA achieves competitive performance compared to standard Multi-Head Attention (MHA), while greatly improving inference speed and GPU memory usage. For example, on a English-German speech translation task, MTLA achieves a 5.3x speedup and a reduction in GPU memory usage by a factor of 8.3 compared to MHA, while maintaining translation quality.

Figures (2)

• Figure 1: Illustration of MTLA. Blue arrows denote transformations by linear layers, and the red dashed lines indicate content attended to during attention. The example corresponds to 4 attention heads. (a) Standard MHA; (b) MTLA with a temporal compression ratio of 2. $\bm{\oplus}$ denotes addition. The transformation from compressed temporal-latent KV cache to multi-head KVs can be absorbed into the query/output linear layers via matrix multiplication associativity, avoiding redundant computation.
• Figure 2: Illustration of MTLA inference and training with temporal compression ratio $2$. $\bm{q}_i$: query, $\bm{x}_i$: attention input, $\hat{\bm{c}}_j$: compressed KV cache, $\hat{\bm{c}}_j^{'}$: temporary version updated later. (a) Incremental inference in MTLA, where at certain steps (e.g., 1, 3, 5), the model attends to the temporary $\hat{\bm{c}}_j^{'}$. (b) KV cache generated by simple pre-downsampling, which mismatches MTLA inference. (c) MTLA training, where a stride-aware causal mask is used to match the inference condition.