Latent Flow Transformer
This work addresses the problem of computational inefficiency in large language models for AI researchers and practitioners, offering a novel compression method that is incremental in improving flow-based approaches.
The paper tackles the inefficiency of discrete layers in transformers by proposing the Latent Flow Transformer (LFT), which compresses multiple layers into a single learned transport operator using flow matching and a Flow Walking algorithm, achieving improved performance with reduced KL divergence (e.g., 0.407 vs. 0.529 for 6-layer compression).
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 \textit{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.