Latent Flow Transformer

Latent Flow Transformer

May 20, 2025

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

cs.AI

Zusammenfassung

Transformer, die Standardimplementierung für große Sprachmodelle (LLMs), bestehen typischerweise aus Dutzenden bis Hunderten von diskreten Schichten. Obwohl mehr Schichten zu einer besseren Leistung führen können, wurde dieser Ansatz als weitgehend ineffizient kritisiert, insbesondere angesichts der Überlegenheit kontinuierlicher Schichten, die durch Diffusions- und flussbasierte Modelle für die Bildgenerierung demonstriert wurde. Wir schlagen den Latent Flow Transformer (LFT) vor, der einen Block von Schichten durch einen einzigen gelernten Transportoperator ersetzt, der über Flow Matching trainiert wird. Dies bietet eine signifikante Kompression bei gleichzeitiger Kompatibilität mit der ursprünglichen Architektur. Zusätzlich adressieren wir die Einschränkungen bestehender flussbasierter Methoden bei der Bewahrung von Kopplungen durch die Einführung des Flow Walking (FW) Algorithmus. Beim Pythia-410M-Modell komprimiert der mit Flow Matching trainierte LFT 6 von 24 Schichten und übertrifft das direkte Überspringen von 2 Schichten (KL-Divergenz der LM-Logits bei 0,407 vs. 0,529), was die Machbarkeit dieses Designs demonstriert. Wenn der LFT mit FW trainiert wird, destilliert er weitere 12 Schichten in eine einzige Schicht, während die KL-Divergenz auf 0,736 reduziert wird, was das Überspringen von 3 Schichten (0,932) übertrifft und die Lücke zwischen autoregressiven und flussbasierten Generierungsparadigmen erheblich verringert.

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.