LaVin-DiT: Large Vision Diffusion Transformer

TL;DR

LaVin-DiT is a scalable vision foundation model using diffusion transformers and a variational autoencoder, enabling multi-task learning and state-of-the-art performance across diverse vision tasks without fine-tuning.

Contribution

The paper introduces LaVin-DiT, a novel large vision diffusion transformer that integrates a variational autoencoder and in-context learning for unified multi-task vision modeling.

Findings

Achieves state-of-the-art results on multiple vision benchmarks.

Scales from 0.1B to 3.4B parameters with effective multi-task performance.

Demonstrates strong generalization without fine-tuning across tasks.

Abstract

This paper presents the Large Vision Diffusion Transformer (LaVin-DiT), a scalable and unified foundation model designed to tackle over 20 computer vision tasks in a generative framework. Unlike existing large vision models directly adapted from natural language processing architectures, which rely on less efficient autoregressive techniques and disrupt spatial relationships essential for vision data, LaVin-DiT introduces key innovations to optimize generative performance for vision tasks. First, to address the high dimensionality of visual data, we incorporate a spatial-temporal variational autoencoder that encodes data into a continuous latent space. Second, for generative modeling, we develop a joint diffusion transformer that progressively produces vision outputs. Third, for unified multi-task training, in-context learning is implemented. Input-target pairs serve as task context, which guides the diffusion transformer to align outputs with specific tasks within the latent space. During inference, a task-specific context set and test data as queries allow LaVin-DiT to generalize across tasks without fine-tuning. Trained on extensive vision datasets, the model is scaled from 0.1B to 3.4B parameters, demonstrating substantial scalability and state-of-the-art performance across diverse vision tasks. This work introduces a novel pathway for large vision foundation models, underscoring the promising potential of diffusion transformers. The code and models are available.