Latente Stroom Transformer

Latent Flow Transformer

May 20, 2025

Auteurs: Yen-Chen Wu, Feng-Ting Liao, Meng-Hsi Chen, Pei-Chen Ho, Farhang Nabiei, Da-shan Shiu

cs.AI

Samenvatting

Transformers, de standaardimplementatie voor grote taalmodellen (LLMs), bestaan doorgaans uit tientallen tot honderden discrete lagen. Hoewel meer lagen tot betere prestaties kunnen leiden, wordt deze aanpak bekritiseerd als verre van efficiënt, vooral gezien de superioriteit van continue lagen die is aangetoond door diffusie- en stroomgebaseerde modellen voor beeldgeneratie. Wij stellen de Latent Flow Transformer (LFT) voor, die een blok lagen vervangt door een enkele geleerde transportoperator die wordt getraind via flow matching, wat aanzienlijke compressie biedt terwijl compatibiliteit met de oorspronkelijke architectuur behouden blijft. Daarnaast pakken we de beperkingen van bestaande stroomgebaseerde methoden in het behouden van koppeling aan door de introductie van het Flow Walking (FW)-algoritme. Op het Pythia-410M-model comprimeert LFT, getraind met flow matching, 6 van de 24 lagen en presteert het beter dan het direct overslaan van 2 lagen (KL-divergentie van LM-logits op 0.407 vs. 0.529), wat de haalbaarheid van dit ontwerp aantoont. Wanneer LFT wordt getraind met FW, destilleert het verder 12 lagen tot één terwijl de KL wordt teruggebracht tot 0.736, wat beter is dan het overslaan van 3 lagen (0.932), waardoor de kloof tussen autoregressieve en stroomgebaseerde generatieparadigma's aanzienlijk wordt verkleind.

English

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