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

TL;DR

LaTtE-Flow is a unified multimodal transformer architecture that enhances image understanding and generation efficiency, achieving faster inference and strong task performance by distributing flow processes across specialized layers.

Contribution

The paper introduces LaTtE-Flow, a novel architecture that combines pretrained VLMs with a layerwise flow-based design for efficient, fast multimodal image understanding and generation.

Findings

Achieves around 6x faster inference speed than recent models.

Maintains strong performance on multimodal understanding tasks.

Provides competitive image generation quality.

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 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.