Lavida-R1: Advancing Reasoning for Unified Multimodal Diffusion Language Models
Abstract
Diffusion language models (dLLMs) recently emerged as a promising alternative to auto-regressive LLMs. The latest works further extended it to multimodal understanding and generation tasks. In this work, we propose LaViDa-R1, a multimodal, general-purpose reasoning dLLM. Unlike existing works that build reasoning dLLMs through task-specific reinforcement learning, LaViDa-R1 incorporates diverse multimodal understanding and generation tasks in a unified manner. In particular, LaViDa-R1 is built with a novel unified post-training framework that seamlessly integrates supervised finetuning (SFT) and multi-task reinforcement learning (RL). It employs several novel training techniques, including answer-forcing, tree search, and complementary likelihood estimation, to enhance effectiveness and scalability. Extensive experiments demonstrate LaViDa-R1's strong performance on a wide range of multimodal tasks, including visual math reasoning, reason-intensive grounding, and image editing.
Lay Summary
Large language models are built to generate text one word at a time. In this work, we explore a newer type of AI model that works differently and can better handle many kinds of information together, such as text and images. We introduce LaViDa-R1, a single generative AI system designed to solve a wide variety of tasks that require reasoning. It can understand images, answer complex visual questions, edit images based on instructions, and solve visual math problems. Instead of training separate systems for different tasks, LaViDa-R1 learns all of them together in one unified framework. To make the system more effective, we develop several new training strategies that help it learn step-by-step reasoning, explore multiple possible solutions, and improve from feedback. These methods allow the model to scale to many different tasks while remaining efficient. Our experiments show that LaViDa-R1 performs strongly across many challenging multimodal tasks, especially those that require deeper reasoning and understanding of both language and images.