Title: LaTtE-Flow: Layerwise Timestep-Expert Flow-based Transformer
Primary area: generative models
Keywords: Diffusion, flow-matching, generative models, efficiency, timestep experts

Abstract: Recent advances in multimodal foundation models capable of both image understanding and generation have opened exciting avenues for building unified systems that seamlessly handle diverse vision-language tasks. Despite the 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 up to 48× faster inference speed compared to existing unified multimodal models.

TL;DR: We propose a novel layerwise timestep-expert flow-based architecture for efficient image generation