Transformador de Fluxo Latente

Resumo

Transformers, a implementação padrão para modelos de linguagem de grande escala (LLMs), geralmente consistem em dezenas a centenas de camadas discretas. Embora mais camadas possam levar a um melhor desempenho, essa abordagem tem sido questionada como pouco eficiente, especialmente considerando a superioridade das camadas contínuas demonstrada por modelos baseados em difusão e fluxo para geração de imagens. Propomos o Latent Flow Transformer (LFT), que substitui um bloco de camadas por um único operador de transporte aprendido, treinado via correspondência de fluxo, oferecendo uma compressão significativa enquanto mantém a compatibilidade com a arquitetura original. Além disso, abordamos as limitações dos métodos baseados em fluxo existentes na preservação do acoplamento, introduzindo o algoritmo Flow Walking (FW). No modelo Pythia-410M, o LFT treinado com correspondência de fluxo comprime 6 de 24 camadas e supera a omissão direta de 2 camadas (Divergência KL dos logits do LM em 0,407 vs. 0,529), demonstrando a viabilidade desse design. Quando treinado com FW, o LFT destila ainda mais 12 camadas em uma, reduzindo a KL para 0,736, superando a omissão de 3 camadas (0,932), estreitando significativamente a lacuna entre os paradigmas de geração autoregressiva e baseada em fluxo.

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.