Latent Flow Transformer

TL;DR

The paper introduces the Latent Flow Transformer, a novel approach that compresses transformer layers using flow matching and flow walking, achieving efficient model compression while maintaining performance.

Contribution

It presents the Latent Flow Transformer, which replaces multiple layers with a learned transport operator and introduces flow walking to improve flow-based model compression.

Findings

LFT compresses 6 of 24 layers, outperforming skipping 2 layers.

LFT with flow walking compresses 12 layers into one, reducing KL divergence.

The method narrows the gap between autoregressive and flow-based models.

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 \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.