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

Ying Shen, Zhiyang Xu, Jiuhai Chen, Shizhe Diao, Jiaxin Zhang, Yuguang Yao, Joy Rimchala, Ismini Lourentzou, Lifu Huang
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:111338-111356, 2026.

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$\times$ faster inference speed compared to recent unified multimodal models.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-shen26w, title = {{L}a{T}t{E}-Flow: Layerwise Timestep-Expert Flow-based Transformer}, author = {Shen, Ying and Xu, Zhiyang and Chen, Jiuhai and Diao, Shizhe and Zhang, Jiaxin and Yao, Yuguang and Rimchala, Joy and Lourentzou, Ismini and Huang, Lifu}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {111338--111356}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/shen26w/shen26w.pdf}, url = {https://proceedings.mlr.press/v306/shen26w.html}, 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$\times$ faster inference speed compared to recent unified multimodal models.} }
Endnote
%0 Conference Paper %T LaTtE-Flow: Layerwise Timestep-Expert Flow-based Transformer %A Ying Shen %A Zhiyang Xu %A Jiuhai Chen %A Shizhe Diao %A Jiaxin Zhang %A Yuguang Yao %A Joy Rimchala %A Ismini Lourentzou %A Lifu Huang %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-shen26w %I PMLR %P 111338--111356 %U https://proceedings.mlr.press/v306/shen26w.html %V 306 %X 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$\times$ faster inference speed compared to recent unified multimodal models.
APA
Shen, Y., Xu, Z., Chen, J., Diao, S., Zhang, J., Yao, Y., Rimchala, J., Lourentzou, I. & Huang, L.. (2026). LaTtE-Flow: Layerwise Timestep-Expert Flow-based Transformer. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:111338-111356 Available from https://proceedings.mlr.press/v306/shen26w.html.

Related Material