mtla

MTLA: Multi-head Temporal Latent Attention

Multi-head Temporal Latent Attention
Keqi Deng, Philip C. Woodland
Paper on arXiv
Accepted at NeurIPS 2025!

About

MTLA is a novel attention mechanism building on DeepSeek MLA, with a key innovation: temporal compression of the key-value cache. This enables more efficient self-attention and significantly reduces memory footprint during inference, making it particularly valuable for decoder-only architectures such as LLMs. Built on PyTorch, this project also serves as an open-source, decoder-only toolkit for end-to-end speech and language processing, covering tasks such as text summarisation, speech translation, speech recognition, spoken language understanding, and so on, with fully featured setup recipes.

Key Features

Supported Attention Mechanisms

• Attention: Multi-head Attention (MHA), Multi-Query Attention (MQA), Grouped-Query Attention (GQA), Multi-head Latent Attention (MLA), and Multi-head Temporal Latent Attention (MTLA)
• Positional Encoding: Rotary Position Embedding (RoPE), and Decoupled Rotary Position Embedding
• FlashAttention: Extended FlashAttention-2 for MTLA inference
• HuggingFace Transformers: Support HuggingFace Transformers toolkit usage to train LLMs based on MTLA

Complete Setup Recipes

• Tasks: speech translation (MuST-C), speech recognition (AMI), spoken language understanding (SLURP), and text summarisation (XSum)
• Data Processing: Fairseq-style Fbank feature extraction and compression into zip file, and ESPnet2-style speech data processing with raw audio saved in flac or ark format
• Feature Extraction: Fbank online/offline extraction, and self-supervised learning representations as features, using upstream models in S3PRL
• Notebook Demo:

Evaluation

• Parallel Inference: Fairseq-style parallel beam search over batches containing multiple data samples
• Quality Evaluation: BLEU, WER, classification accuracy, and ROUGE (ROUGE-1, ROUGE-2, and ROUGE-L)
• Efficiency Evaluation: inference time spent, and GPU memory (including activation memory and the storage of key-value cache) consumed on inference

Installation and Usage

• If you only need the Python MTLA module, simply clone this repository or pip install:

  bash复制

  pip install mtla

  Then refer to the following example:

  python复制

  import torch
  from mtla import MultiheadTemporalLatentAttention
  
  batch, length, dim = 2, 64, 512
  x = torch.randn(batch, length, dim)
  pos = torch.arange(0, length).float().view(1, -1) # Position information
  model = MultiheadTemporalLatentAttention(
      embed_dim=dim, # Model dimension
      num_heads=8,  # Attention heads of queries
  )
  y = model(query=x, key=x, value=x, position=pos)
  assert y.shape == x.shape

  A notebook demo of training with MTLA and performing beam search inference refers to

• Optional: FlashAttention backend for MTLA inference. We provide an optional FlashAttention backend to accelerate MTLA inference. This feature is disabled by default. To enable it, please install our customised FlashAttention fork:

  bash复制

  git clone https://github.com/D-Keqi/flash-attention.git
  cd flash-attention
  python setup.py install

  • FlashAttention requires a CUDA-capable GPU with PyTorch 2.7.0 and CUDA 12.6 (tested working versions).
  • Only fp16 (torch.float16) or bf16 (torch.bfloat16) dtypes are supported.
  • If FlashAttention is not installed, MTLA will automatically fall back to the standard PyTorch implementation.

  Refer to the example below to use our extended FlashAttention for MTLA inference：

…

• If you want to use MTLA through HuggingFace Transformers or train an LLM based on MTLA, you just need to import mtla, then you can load MTLA-based models as easily as you would load any other model in Transformers. See the example below for reference:

…

• If you intend to run the full experiments, please install the project as described below before proceeding to the examples in the experiments directory.
  • PyTorch version >= 1.10.0
  • Python version >= 3.8
  bash复制

  cd experiments/tools/fairseq
  pip install --editable ./

Citation

If you use this codebase, or otherwise find our work valuable, please cite MTLA:

@inproceedings{deng2025mtla,
  title={Multi-head Temporal Latent Attention},
  author={Deng, Keqi and Woodland, Philip C},
  booktitle={Proc. NeurIPS},
  address={San Diego, USA},
  year={2025}
}