Research map / Architectures and efficiency / Attention and transformers

Mini-Sequence Transformers: Optimizing Intermediate Memory for Long Sequences Training

NeurIPS 2024 · Architectures and efficiency › Attention and transformers · approach: Sparse efficient transformers for long context · 1 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes Mini-Sequence Transformer (MsT), which splits MLP and LM-Head computation into sequence-wise mini-sequences to cut intermediate activation memory, enabling 12-24x longer training contexts on a single A100 without throughput or convergence loss.

Problem

Long-sequence LLM training is limited by large intermediate activations in MLP and LM-Head blocks (e.g., logits over large vocabularies), which cap context length even with activation recomputation.

Approach

Partition each input sequence into mini-sequences processed iteratively inside the MLP and LM-Head blocks, with gradient accumulation in the backward pass, so only S/M of intermediate tensors are materialized at once; combined with activation recomputation and chunk-based sizing (C=d) to avoid throughput loss.

Key results

• Llama3-8B with MsT reaches 60K context on a single A100-80GB versus 5K vanilla and 14K with activation recomputation (about 12x vanilla).
• Gemma-2-9B extends maximum sequence length 24x over vanilla (1.5K to 36K), Qwen2-7B 18x, Mistral-7B 12x.
• MsT keeps throughput close to activation recomputation (Llama3-8B at 8K: 3194.90 vs 3271.42 TFLOPS at batch size 2), and LongAlpaca training at 8K context gives loss 2.00 vs 2.23 for activation recomputation.

Measured effects reported in the paper

Surprising findings

• Training Llama3-8B at 30K context gave perplexity 3.45 versus 7.41 at 8K under MsT, but the gap largely reflects the longer context rather than a method-to-method comparison.
• Combining MsT with 4-bit quantization reaches 140K tokens, though the authors did not evaluate the effect of quantization on training loss.

Limitations

• Built on PyTorch rather than fused CUDA kernels, which may limit performance and memory savings.
• Throughput drops for short sequences where M-fold weight-gradient overhead dominates, requiring chunk-based mitigation.
• Quantization combinations were not evaluated for training loss.

Experimental setup

Datasets: LongAlpaca-12k

Benchmarks: LongAlpaca-12k

Metrics: maximum sequence length training throughput (TFLOPS) training time per step perplexity loss

Baselines: vanilla PyTorch / HuggingFace training activation recomputation 8-bit quantization 4-bit quantization Huggingface PEFT

Builds on: activation recomputation FlashAttention-2 DeepSpeed-Ulysses HuggingFace Transformers

Closest papers by meaning

• StreamBP: Memory-Efficient Exact Backpropagation for Long Sequence Training of LLMs (NeurIPS 2025 · similarity 0.88)
• Megalodon: Efficient LLM Pretraining and Inference with Unlimited Context Length (NeurIPS 2024 · similarity 0.87)
• MISA: Memory-Efficient LLMs Optimization with Module-wise Importance Sampling (NeurIPS 2025 · similarity 0.86)
• IceFormer: Accelerated Inference with Long-Sequence Transformers on CPUs (ICLR 2024 · similarity 0.86)
• ShiftAddLLM: Accelerating Pretrained LLMs via Post-Training Multiplication-Less Reparameterization (NeurIPS 2024 · similarity 0.86)
• Fira: Can We Achieve Full-rank Training of LLMs Under Low-rank Constraint? (NeurIPS 2025 · similarity 0.86)