LaTtE-Flow: Layerwise Timestep-Expert Flow-based Transformer
Abstract
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 often rely on extensive pretraining and suffer from slow generation speeds, limiting their practical deployment in real-time and resource-constrained settings. In this work, we introduce Layerwise Timestep-Expert Flow-based Transformer (LaTtE-Flow), a novel architecture that improves the efficiency of diffusion/flow-based Transformers within the unified model setting. LaTtE-Flow builds upon powerful pre-trained Vision-Language Models (VLMs) to inherit strong multimodal understanding capabilities, and extends them with a novel Layer-wise 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 6× faster inference speed compared to recent unified multimodal models.
Lay Summary
Unified vision-language models can both understand images and text and generate new images, but they are often too slow for practical use. Image generation usually requires many sampling steps, and at each step, the model needs to run through a large number of Transformer layers. This creates a lot of repeated computation, making generation expensive and difficult to deploy in real-time or resource-limited settings. LaTtE-Flow is designed to make this process more efficient. Instead of using the whole model at every generation step, it divides Transformer layers into specialized groups, where each group handles a different range of timesteps in the flow-matching process. During sampling, only the layers needed for the current timestep are activated, so the model avoids wasting computation on layers that are less useful at that stage. This layerwise specialization allows LaTtE-Flow to generate images much faster while still preserving strong multimodal understanding from pretrained vision-language models. LaTtE-Flow reports around 6× faster inference compared with recent unified multimodal models, while maintaining competitive image generation quality. This efficiency improvement is important because it makes unified models more practical for applications that need fast responses or have limited computing resources.