Transformateur de Flux Latent

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

Résumé

Les Transformers, l'implémentation standard pour les grands modèles de langage (LLMs), sont généralement composés de dizaines à centaines de couches discrètes. Bien qu'un plus grand nombre de couches puisse améliorer les performances, cette approche a été remise en question comme étant loin d'être efficace, surtout compte tenu de la supériorité des couches continues démontrée par les modèles de diffusion et basés sur les flux pour la génération d'images. Nous proposons le Latent Flow Transformer (LFT), qui remplace un bloc de couches par un unique opérateur de transport appris via le flow matching, offrant une compression significative tout en restant compatible avec l'architecture originale. De plus, nous abordons les limitations des méthodes existantes basées sur les flux en matière de préservation du couplage en introduisant l'algorithme Flow Walking (FW). Sur le modèle Pythia-410M, le LFT entraîné avec le flow matching compresse 6 des 24 couches et surpasse le saut direct de 2 couches (Divergence KL des logits du modèle de langage à 0.407 contre 0.529), démontrant la faisabilité de cette conception. Lorsqu'il est entraîné avec FW, le LFT distille en outre 12 couches en une tout en réduisant la KL à 0.736, surpassant celle obtenue en sautant 3 couches (0.932), réduisant ainsi significativement l'écart entre les paradigmes de génération autoregressive et basée sur les flux.

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.