Latent Flow Transformer
Journal:
arXiv
Published Date:
May 20, 2025
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.
