Latent Flow Transformer

2025-05-20

AI Summary

To address layer redundancy and inefficiency arising from discrete architectural designs in large language models (LLMs), this paper proposes the Latent Flow Transformer (LFT), which replaces multiple discrete Transformer blocks with a learnable, continuous transport operator—enabling structural compression while preserving compatibility with the original architecture. We introduce flow matching—a first-of-its-kind application in LLM compression—and propose the Flow Walking algorithm to mitigate coupling degradation in streaming generation, all while retaining autoregressive modeling capability. On Pythia-410M, LFT achieves a KL divergence of 0.407 after compressing 6 layers (outperforming layer-skipping baselines); when compressing 12 layers into 1, KL divergence is 0.736—significantly narrowing the performance gap between autoregressive and streaming generation paradigms.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Transformers, the standard implementation for large language models (LLMs), typically consist of tens to hundreds of discrete layers. While more layers can lead to better performance, this approach has been challenged as far from efficient, especially given the superiority of continuous layers demonstrated by diffusion and flow-based models for image generation. We propose the Latent Flow Transformer (LFT), which replaces a block of layers with a single learned transport operator trained via flow matching, offering significant compression while maintaining compatibility with the original architecture. Additionally, we address the limitations of existing flow-based methods in extit{preserving coupling} by introducing the Flow Walking (FW) algorithm. On the Pythia-410M model, LFT trained with flow matching compresses 6 of 24 layers and outperforms directly skipping 2 layers (KL Divergence of LM logits at 0.407 vs. 0.529), demonstrating the feasibility of this design. When trained with FW, LFT further distills 12 layers into one while reducing the KL to 0.736 surpassing that from skipping 3 layers (0.932), significantly narrowing the gap between autoregressive and flow-based generation paradigms.

Problem

• Replacing discrete transformer layers with continuous flow-based operators
• Improving efficiency by compressing multiple layers into one
• Enhancing flow-based methods to preserve coupling in text generation

Innovation

• Replaces layers with learned transport operator
• Introduces Flow Walking algorithm for coupling
• Compresses layers while maintaining performance