Research map / Architectures and efficiency / Attention and transformers

FFN Fusion: Rethinking Sequential Computation in Large Language Models

NeurIPS 2025 · Architectures and efficiency › Attention and transformers · approach: Efficient language transformers · 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 FFN Fusion, which parallelizes runs of consecutive FFN layers left after attention removal (via Puzzle NAS) by summing them into one wider FFN, and applies it to derive the 253B Ultra-253B-Base from Llama-3.1-405B-Instruct.

Problem

Sequential depth in large transformer LLMs creates latency bottlenecks under tensor parallelism (synchronization and small-GEMM overheads), and existing compression methods (quantization, pruning, MoE) do not directly reduce sequential computation.

Approach

Consecutive FFN layers remaining after attention removal have low inter-layer dependency, so they can be applied in parallel to the same input and exactly merged into a single wider FFN by concatenating SwiGLU weight matrices; the paper also probes pairwise block dependency to find parallelizable full blocks.

Key results

• Ultra-253B-Base reaches 1.71x user-latency speedup over Llama-405B (70.92 vs 41.44 tokens/s, TP8 on H100) and 35x lower per-token cost at batch size 32.
• Ultra-253B-Base scores 84.92 on Arena Hard vs 72.56 for Llama-405B, 87.54 vs 86.53 on MMLU Instruct, and 9.19 vs 9.06 on MT-Bench.
• On Puzzle-49B, step-3 FFN fusion gives 10% latency gain for 3.5% accuracy drop; a 25B-token KD phase lifts the fused model to MMLU 80.56 and MT-Bench 9.00, above the baseline.

Measured effects reported in the paper

Surprising findings

• Fusing all 49 consecutive FFN layers after attention removal keeps MMLU at 82.76 (vs 84.23 before), even though removing those layers would severely damage performance.
• Full transformer blocks (attention plus FFN) can be parallelized in small groups with modest MMLU loss, though the drop becomes steep at a third group.

Limitations

• Fusion is constrained by per-GPU memory, requiring splitting into several fusion sequences.
• Full-block parallelization degrades sharply beyond two groups and is not supported natively by TensorRT-LLM or vLLM; experiments used HuggingFace Transformers with reduced optimization.
• Final FFN in each sequence is sensitive to fusion and typically must be excluded.

Experimental setup

Datasets: Distillation Mix FineWeb Dolma Buzz-V1.2

Benchmarks: MMLU MMLU-Pro Arena Hard HumanEval MT-Bench

Metrics: accuracy tokens/second speedup per-token cost

Baselines: Llama-3.1-405B-Instruct Llama-3.3-70B Puzzle-49B Ultra-PreFusion

Builds on: Puzzle SwiGLU speculative decoding

Closest papers by meaning

• Deep Fusion: Efficient Network Training via Pre-trained Initializations (ICML 2024 · similarity 0.88)
• Skip a Layer or Loop It? Learning Program-of-Layers in LLMs (ICML 2026 · similarity 0.88)
• Search for Efficient Large Language Models (NeurIPS 2024 · similarity 0.88)
• Memory Layers at Scale (ICML 2025 · similarity 0.87)
• Layer as Puzzle Pieces: Compressing Large Language Models through Layer Concatenation (NeurIPS 2025 · similarity 0.87)
• Polar Sparsity: High Throughput Batched LLM Inferencing with Scalable Contextual Sparsity (NeurIPS 2025 · similarity 0.87)