VIRAL: Visual In-Context Reasoning via Analogy in Diffusion Transformers

TL;DR

The paper tackles visual in-context learning (V-ICL) for diverse vision tasks by enabling task-agnostic inference from a single demonstration. It introduces VIRAL, recasting V-ICL as visual analogy conditional generation on a frozen Diffusion Transformer with role-aware multi-image conditioning and Mixture-of-Experts LoRA. A large-scale In-Context Editing Dataset spanning perception, restoration, and open-domain editing supports learning from exemplar-query quadruplets. Empirical results demonstrate strong V-ICL performance, competitive results with task-specific models, and evidence that the visual analogy formulation $x_s : x_t :: x_q : y_q$ enables a unified visual generalist.

Abstract

Replicating In-Context Learning (ICL) in computer vision remains challenging due to task heterogeneity. We propose \textbf{VIRAL}, a framework that elicits visual reasoning from a pre-trained image editing model by formulating ICL as conditional generation via visual analogy ($x_s : x_t :: x_q : y_q$). We adapt a frozen Diffusion Transformer (DiT) using role-aware multi-image conditioning and introduce a Mixture-of-Experts LoRA to mitigate gradient interference across diverse tasks. Additionally, to bridge the gaps in current visual context datasets, we curate a large-scale dataset spanning perception, restoration, and editing. Experiments demonstrate that VIRAL outperforms existing methods, validating that a unified V-ICL paradigm can handle the majority of visual tasks, including open-domain editing. Our code is available at https://anonymous.4open.science/r/VIRAL-744A

Figures (18)

• Figure 1: Illustration of in-context learning with VIRAL. Given a reference exemplar pair, VIRAL interprets the underlying visual transformation and applies it to a query image, including standard visual task and open-domain editing.
• Figure 2: Overview of the proposed Visual In-Context Learning framework. We unify diverse visual tasks into a homogeneous RGB pixel space, enabling a universal generative interface. The visual tokens from the reference exemplar pair and the query image are concatenated along the sequence dimension and fed into the Diffusion Transformer (DiT) backbone. The exemplar pair and the query images remain fixed during denoising steps, while the model updates only the noisy latent to the target image.
• Figure 3: Quantitative comparison. We evaluate the performance of four V-ICL baselines against VIRAL. While existing baselines either exhibit restricted task versatility or suffer from performance degradation when encountering complex scenarios due to their reliance on over-simplified training distributions, our model consistently achieves superior accuracy and visual fidelity across all evaluated tasks.
• Figure 4: Quantitative comparison on open-domain editing tasks. For style transfer, instruction-driven models often struggle to achieve the desired results. In contrast, our visual in-context demonstrations ensure both stylistic consistency and content preservation. The text editing instructions are provided in Appendix \ref{['sec:appendix_prompts']}.
• Figure 5: Zero-shot generalization to unseen Lineart Generation. Despite being trained exclusively on standard Canny edge maps, VIRAL successfully generalizes to the artistic lineart task via one-shot in-context learning.