LaTtE-Flow: Transformador Baseado em Fluxo com Especialistas por Camada e Timestep

LaTtE-Flow: Layerwise Timestep-Expert Flow-based Transformer

June 8, 2025

Autores: Ying Shen, Zhiyang Xu, Jiuhai Chen, Shizhe Diao, Jiaxin Zhang, Yuguang Yao, Joy Rimchala, Ismini Lourentzou, Lifu Huang

cs.AI

Resumo

Avanços recentes em modelos de base multimodal que unificam a compreensão e geração de imagens abriram caminhos promissores para abordar uma ampla gama de tarefas de visão e linguagem dentro de um único framework. Apesar do progresso, os modelos unificados existentes geralmente exigem um extenso pré-treinamento e lutam para alcançar o mesmo nível de desempenho em comparação com modelos dedicados a cada tarefa. Além disso, muitos desses modelos sofrem com velocidades lentas de geração de imagens, limitando sua implantação prática em cenários em tempo real ou com recursos limitados. Neste trabalho, propomos o Layerwise Timestep-Expert Flow-based Transformer (LaTtE-Flow), uma arquitetura nova e eficiente que unifica a compreensão e geração de imagens dentro de um único modelo multimodal. O LaTtE-Flow se baseia em poderosos modelos de Visão e Linguagem (VLMs) pré-treinados para herdar capacidades robustas de compreensão multimodal, e os estende com uma nova arquitetura baseada em fluxo de especialistas por camadas e timesteps para geração eficiente de imagens. O LaTtE-Flow distribui o processo de correspondência de fluxo entre grupos especializados de camadas Transformer, cada um responsável por um subconjunto distinto de timesteps. Esse design melhora significativamente a eficiência de amostragem ao ativar apenas um pequeno subconjunto de camadas em cada timestep de amostragem. Para aprimorar ainda mais o desempenho, propomos um mecanismo de Atenção Residual Condicionada por Timestep para reutilização eficiente de informações entre camadas. Experimentos demonstram que o LaTtE-Flow alcança um desempenho robusto em tarefas de compreensão multimodal, enquanto obtém qualidade competitiva na geração de imagens com uma velocidade de inferência aproximadamente 6x mais rápida em comparação com modelos multimodais unificados recentes.

English

Recent advances in multimodal foundation models unifying image understanding and generation have opened exciting avenues for tackling a wide range of vision-language tasks within a single framework. Despite progress, existing unified models typically require extensive pretraining and struggle to achieve the same level of performance compared to models dedicated to each task. Additionally, many of these models suffer from slow image generation speeds, limiting their practical deployment in real-time or resource-constrained settings. In this work, we propose Layerwise Timestep-Expert Flow-based Transformer (LaTtE-Flow), a novel and efficient architecture that unifies image understanding and generation within a single multimodal model. LaTtE-Flow builds upon powerful pretrained Vision-Language Models (VLMs) to inherit strong multimodal understanding capabilities, and extends them with a novel Layerwise Timestep Experts flow-based architecture for efficient image generation. LaTtE-Flow distributes the flow-matching process across specialized groups of Transformer layers, each responsible for a distinct subset of timesteps. This design significantly improves sampling efficiency by activating only a small subset of layers at each sampling timestep. To further enhance performance, we propose a Timestep-Conditioned Residual Attention mechanism for efficient information reuse across layers. Experiments demonstrate that LaTtE-Flow achieves strong performance on multimodal understanding tasks, while achieving competitive image generation quality with around 6x faster inference speed compared to recent unified multimodal models.