Transformador de Flujo Latente

Latent Flow Transformer

May 20, 2025

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

cs.AI

Resumen

Los Transformers, la implementación estándar para los modelos de lenguaje de gran escala (LLMs), suelen consistir en decenas o cientos de capas discretas. Aunque un mayor número de capas puede mejorar el rendimiento, este enfoque ha sido cuestionado por su falta de eficiencia, especialmente dada la superioridad de las capas continuas demostrada por los modelos basados en difusión y flujo para la generación de imágenes. Proponemos el Transformer de Flujo Latente (LFT), que reemplaza un bloque de capas con un único operador de transporte aprendido mediante el emparejamiento de flujos, ofreciendo una compresión significativa mientras mantiene la compatibilidad con la arquitectura original. Además, abordamos las limitaciones de los métodos basados en flujo existentes para preservar el acoplamiento mediante la introducción del algoritmo Caminata de Flujo (FW). En el modelo Pythia-410M, el LFT entrenado con emparejamiento de flujos comprime 6 de 24 capas y supera el rendimiento de saltar directamente 2 capas (Divergencia KL de los logits del modelo de lenguaje en 0.407 frente a 0.529), demostrando la viabilidad de este diseño. Cuando se entrena con FW, el LFT destila además 12 capas en una mientras reduce la KL a 0.736, superando el resultado de saltar 3 capas (0.932), lo que reduce significativamente la brecha entre los paradigmas de generación autoregresiva y basada en flujo.

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.