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.