NeurIPS 2026 Papers

Conference on Neural Information Processing Systems

Decision breakdown

6,437papers

Counted from the official paper list; may differ slightly from the number the conference announced.

  • Oral
    991.5%
  • Spotlight
    2403.7%
  • Poster
    6,09894.7%

Browse all papers

Driving on Memory

Christian Löwens, Thorben Funke, Alexandru P P Condurache

Expand abstract

End-to-end autonomous driving models learn to plan the vehicle's future trajectory from raw sensor observations. While earlier driving benchmarks often measured deviation from the human trajectory, current benchmarks such as NAVSIM and Bench2Drive evaluate models with richer simulation-based metrics intended to capture safe and compliant driving. A high benchmark score should reflect that a model can understand the scene in front of it and act accordingly. But how much of that score specifically comes from reacting to the dynamic part of that scene? To probe this, we remove a model's camera input and replace it with memories from prior drives of the same location. The retrieved memories can provide persistent scene information, including road layout and location-conditioned regularities, but because they are historical, they do not reveal the current actors, signal states, or transient conditions. Surprisingly, memory is nearly sufficient on NAVSIM, reaching or even exceeding the performance of leading end-to-end methods without actually observing the evaluated scene. Our results suggest that a high NAVSIM score does not require a planner to react to the current traffic scene and should therefore be treated with caution. However, we do not observe the same behavior on Bench2Drive: driving from memory remains far from camera-based performance, suggesting that Bench2Drive better preserves interaction demands through longer closed-loop routes. To ensure reproducibility, we provide our code at anonymous.4open.science/r/MemoryDrivoR-E64C.

Poster
Driving Video Retrieval for Complex Queries with Structured Grounding

Manyi Yao, Sparsh Garg, Christian Shelton, Amit Roy-Chowdhury, Abhishek Aich

Expand abstract

Video retrieval at scale is central to data curation and safety validation in autonomous driving, where users want to find not only scenes but also dynamic events such as cut-ins and hard braking. Existing vision-language and keyword-based retrieval methods often miss these events because the relevant motion may not be explicitly described in text or captured by lexical overlap. Rule-based retrieval can encode such events more directly, but it is brittle: generated or hand-written rules often fail when their assumptions do not match real driving data. We propose STRIVE-D, a data-calibrated retrieval framework for driving videos. STRIVE-D uses weakly labeled in-domain videos to estimate when a query rule is reliable, adapt rules that mismatch observed data, and fuse calibrated rule scores with vision-language and keyword-based retrieval signals. Across three driving benchmarks, including newly released human annotations for event retrieval on DrivingDojo, STRIVE-D achieves up to a 30.4% relative improvement over state-of-the-art retrieval models.

Poster
Driver Attention as Competitive Allocation: A Scene--Task-Aware Dual-Branch Framework

Qianfang Wang, Bin Rao, Pengpeng Xu, Tiantian CHEN

Expand abstract

Driver attention prediction is crucial for interpretable driver monitoring and human-like autonomous driving. Existing methods typically formulate this task as a scene-to-heatmap regression problem, failing to reveal the underlying dynamic attention redistribution mechanism. To address this issue, we propose a scene-task-aware competitive allocation framework, which decomposes driver attention as a bottom-up visual saliency prior and a top-down scene-task utility representation. An adaptive product-of-experts allocation rule is designed to fuse these two heterogeneous branches in a competitive manner, where scene-task demand dynamically modulates the dependence on top-down task utility, and driving control urgency adjusts spatial attention concentration via a learnable attention budget temperature. To explicitly differentiate task-critical regions from visually prominent yet task-irrelevant distractors, we incorporate competition-aware contrastive learning to enhance feature discrimination in the latent space. Extensive experiments demonstrate that our method achieves strong and balanced performance, with a Kullback–Leibler divergence of 0.98 and a normalized scanpath saliency of 4.15 on the DR(eye)VE dataset, as well as a Pearson’s correlation coefficient of 0.67 and a histogram similarity of 0.54 on the BDD-A dataset. Ablation studies and diagnostic analyses further validate that our framework yields interpretable attention allocation behaviors, including urgency-induced spatial entropy contraction and scene-task-demand-guided dynamic attention reallocation.

Poster
DriveWAM: Video Generative Priors Enable Scalable World-Action Modeling for Autonomous Driving

Chen Shi, Jinrui Xu, Shaoshuai Shi, Kehua Sheng, Bo Zhang, Li Jiang

Expand abstract

Pretrained foundation models have become an important basis for end-to-end autonomous driving. In contrast to vision-language models pretrained primarily on static image-text pairs, video generative models capture temporal dynamics and motion priors that are naturally suited for driving. We present DriveWAM, a driving world-action model that adapts a pretrained video diffusion transformer into an autoregressive video-action policy. DriveWAM organizes video and action streams into a unified temporal token sequence and trains them under a joint flow-matching objective, preserving the pretrained video-generation architecture while adapting its large-scale video priors to action generation. To incorporate high-level scene understanding, we introduce scene-evolving driving guidance, where a frozen VLM produces chunk-specific semantic intent to guide video-action generation. To keep long-horizon rollout bounded, we further introduce selective KV memory, which maintains bounded modality-aware video and action memory pools through relevance-redundancy cache selection at inference time. Experiments on NAVSIM and the PhysicalAI-Autonomous-Vehicles benchmark show that DriveWAM achieves strong planning performance, and a data-scaling study from 4k to 100k driving clips further confirms the scaling potential of world-action modeling for end-to-end autonomous driving.

Poster
DriveStreamBench: Evaluating User-Conditioned Watch-and-Notify in Streaming Driving Video

Yi Wang, Xin Zhao, Tian Meng, Rui Dai, Chi Zhang, Cui Tang, Bingzheng Liu, Fucheng Wang, Haopeng Zhang, Ruidong Ding, Yang Li, Kaikui Liu, Xiangxiang Chu

Expand abstract

Vision-language models (VLMs) have made steady progress as in-vehicle driving assistants, yet existing evaluations largely remain offline, focusing on perception and short-term reasoning over pre-recorded driving scenes. A complementary capability is still underexplored: whether a model can continuously monitor a causal driving video stream and issue a calibrated alert when a user-specified condition is met. Existing driving benchmarks mainly assess reactive understanding, while recent streaming video benchmarks are developed largely outside the driving domain. To this end, we introduce DriveStreamBench, a driving and streaming video benchmark that organizes this capability spectrum into four cognitive levels constructed from a shared pool of driving events, to support level-wise capability profiling and gated evaluation. We evaluate various vision-language models along with DriveStream-SFT, a supervised reference baseline trained on the benchmark's training split, under a unified streaming protocol. DriveStreamBench tests a basic but necessary capability, and all evaluated models fall short of it; current vision-language models are not yet ready for user-conditioned alerting in driving. Resources are available at https://anonymous.4open.science/r/DriveStreamBench/ and https://dataverse.harvard.edu/previewurl.xhtml?token=01f6bc4c-d809-4e40-b025-b6d783c86571.

Poster
DriveSpatial: A Benchmark for Spatiotemporal Intelligence in VLMs for Autonomous Driving

Anh Hao Vo, Khoa Vo, Phu Loc Nguyen, Sieu Tran, Duc Nguyen, Cuong Ngo, Gladys Gawugah, Sreevenkata A Godavarthi, Chase Rainwater, Nghi Bui, Anh Nguyen, Duy M. H. Nguyen, Ngan Le

Expand abstract

Spatiotemporal intelligence in autonomous driving (AD) requires an agent to integrate multi-view observations into a coherent scene representation, maintain object continuity across viewpoints and time, and reason about spatial relations, interactions, and future dynamics. However, existing AD vision-language benchmarks largely focus on single-view, static, ego-centric, or single-source question answering, leaving it unclear whether current Vision-Language Models (VLMs) can truly construct and reason over dynamic driving scenes. We introduce DriveSpatial a benchmark of 15.6K human-verified QA pairs across 20 tasks from five large-scale AD datasets. DriveSpatial evaluates four abilities: Cognitive Scene Construction, Multi-view Relational Understanding, Temporal Reasoning, and Generalization. Unlike prior benchmarks, DriveSpatial is generated from a dynamic multi-relational scene graph that encodes object states, spatial relations, interactions, camera visibility, and temporal correspondences, enabling QA pairs that enforce genuine cross-view and spatiotemporal reasoning. Evaluating 15 representative VLMs reveals a substantial human–model gap: the strongest model trails humans by 28.4 points, with Cognitive Scene Construction emerging as the key bottleneck. Further diagnostics show that language-only prompting is insufficient, while explicit BEV grounding consistently improves performance. These results suggest that current VLMs lack the scene-construction ability needed for reliable spatiotemporal driving intelligence. DriveSpatial and its construction pipeline will be released to support future research.

Poster
DriveMind: Mind-Evolving Belief Tracking for Closed-Loop Autonomous Driving

Yan Shi, Kehua Chen, Jiarui Qi, Yihong Tang, Yinhai Wang, Liao Haicheng

Expand abstract

Recent multimodal large language models (MLLMs) can perceive a traffic scene, reason about it, and propose a driving action within a single model. However, existing MLLM-based drivers reason from scratch at each frame and commit to a single explanation, discarding earlier hypotheses about the scene. When the scene is ambiguous, the agent locks in on the wrong interpretation and acts on it before evidence accumulates. We propose DriveMind, an MLLM-based driving framework that keeps several hypotheses about the scene alive at the same time and updates their weights as new frames arrive. Inspired by particle filtering, DriveMind propagates each hypothesis through a VLM, scores it against the current observation, and resamples to maintain diversity. On Bench2Drive, DriveMind raises the Driving Score (DS) from 71.36 to 81.37 and the Success Rate (SR) from 45.45% to 54.09%, with the largest gains on ambiguous scenarios such as unsignalized intersections (+20.10 DS) and oncoming-traffic interactions (+20.50 DS). Beyond the overall metrics, DriveMind also improves fine-grained interactive driving abilities such as merging, emergency braking, and give-way. The code will be released soon.

Poster
DriveFuture: Future-Aware Latent World Models for Autonomous Driving

Yufeng hong, Xiaotian Zhou, Yingyan Li, Xiangpo Zhou, Lin Liu, Yadan Luo, Zheng Liu, Lei Yang, Ziying Song

Expand abstract

Existing latent world models for autonomous driving have opened a promising path toward future-aware driving intelligence. However, they typically treat future latent states as prediction targets or auxiliary signals, rather than directly conditioning trajectory planning. This can entangle current and future features in latent space. In this work, we propose DriveFuture, a future-aware latent world modeling framework for autonomous driving that explicitly learns planning-oriented foresight by conditioning the current latent state modeling process on future world states. Specifically, during training, the model first predicts future latent world states from the current latent state and ego action, and then refines the prediction against the ground-truth future latent state via cross-attention. The resulting future-aware latent serves as an explicit condition for a diffusion-based trajectory planner. During inference, DriveFuture conditions on the predicted future latent state instead of the ground-truth future state. DriveFuture achieves SOTA performance on the public NAVSIM benchmarks, reaching 55.5 EPDMS on NAVSIM-v2 navhard, 89.9 EPDMS on NAVSIM-v2 navtest, and 90.7 PDMS on NAVSIM-v1 navtest, respectively. These results suggest that the key to latent world modeling lies not merely in simulating future states, but more importantly in conditioning current decision-making on future states. Notably, as of April 2026, DriveFuture ranks 1st on the NAVSIM-v2 navhard leaderboard and achieves SOTA performance on NAVSIM-v1 navtest.

Poster
DriveDreamer-Policy: A Geometry-Grounded World-Action Model for Unified Generation and Planning

Yang Zhou, Xiaofeng Wang, Hao Shao, Letian Wang, Guosheng Zhao, Jiangnan Shao, Jiagang Zhu, Tingdong Yu, Zheng Zhu, Guan Huang, Steven Waslander

Expand abstract

Recently, world-action models (WAM) have emerged to bridge vision-language-action (VLA) models and world models, unifying their reasoning and instruction-following capabilities and spatio-temporal world modeling. However, existing WAM approaches often focus on modeling 2D appearance or latent representations, with limited geometric grounding, an essential element for embodied systems operating in the physical world. We present DriveDreamer-Policy, a unified driving world-action model that integrates depth generation, future video generation, and motion planning within a single modular architecture. The model employs a large language model to process language instructions, multi-view images, and actions, followed by three lightweight generators that produce depth, future video, and actions. By learning a geometry-aware world representation and using it to guide both future prediction and planning within a unified framework, the proposed model produces more coherent imagined futures and more informed driving actions, while maintaining modularity and controllable cost. Experiments on the Navsim v1 and v2 benchmarks demonstrate that DriveDreamer-Policy achieves strong performance on both planning and world generation tasks. In particular, our model reaches 89.2 PDMS on Navsim v1 and 88.7 EPDMS on Navsim v2, outperforming existing world-model-based approaches while producing higher-quality future video and depth predictions. Ablation studies further show that depth learning provides complementary benefits to video imagination and improves planning robustness. Our code is available to facilitate further research.

Poster
Drive vs. Decay: On the Training Dynamics of Joint-Embedding Predictive Architectures

José Lucas De Melo Costa, Seong Woo Ahn, Fabrice Popineau, Arpad Rimmel, Bich-Liên DOAN

Expand abstract

Joint-Embedding Predictive Architectures (JEPAs) are prone to representation collapse, typically mitigated through empirical heuristics. We develop an early-training stability theory that unifies these heuristics. Linearising the coupled JEPA gradient flow around the trivial fixed point reveals two competing effects: a driving force ($\gamma$) and a decay effect ($\sigma$). Under approximate spectral decoupling, a per-mode stability ratio $\mu_i = \gamma_i / \sigma_i$ factorises into independent data-side and predictor-side terms and the count of unstable modes tracks the rank of representations that can emerge. The framework predicts a phase boundary, which we confirm empirically across more than 800 Tabular-JEPA configurations. It also unifies predictor scaling, masking ratio, and EMA as distinct mechanisms for shifting $\mu$. Guided by this analysis, we introduce ResidualPred, a transformer predictor whose attention is biased toward the identity at initialisation; it improves both effective rank and downstream accuracy on tabular benchmarks and on I-JEPA pretraining over CIFAR-10, CIFAR-100, and STL-10. Our framework connects empirical collapse-avoidance heuristics to an explicit dynamical picture, yielding theory-driven stabilizers. Code is available as supplementary material.

Poster
Drifting Fields are not Conservative

Leonard T. Franz, Sebastian Hoffmann, Tim Weiland, Bernhard Scholkopf, Georg Martius

Expand abstract

Drifting models have recently gained attention for generating high-quality samples in a single forward pass. During training, they learn a push-forward map by following a vector-valued field, the *drift field*. We ask whether this procedure is equivalent to optimizing a scalar loss and find that, in general, it is not: drift fields are *not conservative* and cannot be written as the gradient of any scalar potential. We identify the position-dependent normalization as the source of non-conservatism, with the Gaussian kernel as the unique radial exception. Guided by this, we introduce the *sharp kernel* $k^\\#$ and a sharp-normalized drift field that is conservative for general radial kernels. The resulting vector field is the gradient of a scalar potential that can be optimized directly using stochastic gradient descent. Moreover, the field has the form of a score difference of kernel density estimates, and gives exact equilibrium identifiability. Thus, sharp normalization closes the gap to related literature, such as Wasserstein gradient-flows and denoising score matching, also for non-Gaussian kernels. Empirically, sharp normalization preserves the performance of the original drifting objective, suggesting that the non-conservative flexibility is not required for high-quality generation.

Poster
Drifting Field Policy: Wasserstein Gradient Flow on Policy Space for Offline-to-Online RL

Juil Koo, Mingue Park, Jiwon Choi, Yunhong Min, Minhyuk Sung

Expand abstract

We propose Drifting Field Policy (DFP), a non-ODE one-step generative policy built on the drifting model paradigm. Under online RL finetuning, we frame the policy update as a reverse-KL Wasserstein-2 gradient flow toward a soft target policy, realized by DFP as a steepest-descent step on probability space. By construction, this gradient step is decomposed into an ascent toward higher action-value regions and a score matching with the anchor policy as a trust region. We further derive a simple, tractable surrogate of the otherwise intractable update loss, akin to behavior cloning on top-K critic-selected actions. We find empirically that this mechanism uniquely benefits the drifting backbone owing to its non-ODE parameterization. With one-step inference, DFP achieves state-of-the-art performance on several manipulation tasks across Robomimic and OGBench, outperforming ODE-based policies.

Poster
DriftWeight: Repulsive Drift in Mean-Flow Space for Neural Network Weight Generation

Reo Iizuka, shioya hiroaki, Naoto Yokoya

Expand abstract

Generating neural network weights in a single forward pass promises to amortize model training for fast sampling, transfer, and model-set construction. Yet accuracy alone is insufficient: existing generators can map diverse latent inputs to weights that remain functionally close to their training checkpoints. We study this functional mode collapse in the one-step regime and introduce DriftWeight, a MeanFlow-based generator with temperature-scaled repulsive drift in a checkpoint-augmented PCA proxy space. The method uses structured negatives from the current batch and optimization history to push generated weights away from observed clusters while preserving task performance. Our key finding is that repulsion is not automatically compatible with one-step inference. In a controlled ablation, applying the drift loss at the same boundary evaluation used for sampling collapses one-step accuracy even though multi-step integration remains accurate. DriftWeight mitigates this boundary-localized failure with a consistency objective that anchors the inference point to the training-weight manifold. Across MNIST, Fashion-MNIST, and CIFAR-10, DriftWeight matches or exceeds reported 100-step DeepWeightFlow accuracy in a single forward pass. On MNIST, where we directly evaluate error-set functional overlap, it reduces MaxIoU from the DeepWeightFlow regime of approximately $0.82$ to $0.66$. On CIFAR-100 with a ViT-Base target ($\sim$86M parameters), one-step samples remain within $1.7$ percentage points of the training-population mean.

Poster
Drift-Resistant Navigation World Model with Anchored Epipolar Guidance

Po-Chien Luan, Zimin Xia, Wuyang Li, Yang Gao, Alexandre Alahi

Expand abstract

We propose Drift-Resistant Navigation World Model, a generative model that mitigates both perceptual drift and geometric drift in conventional rollout-based navigation world models. Existing methods recursively feed generated content into subsequent steps, causing noise accumulation and degraded predictions, i.e., perceptual drift. Meanwhile, their predictions often deviate from the agent’s motion, resulting in geometry drift. We address both types of drift by redesigning world-model prediction as an anchor-guided rollout. Instead of rolling out every frame sequentially, we first predict sparse future anchors that serve as stable long-range targets, and then generate intermediate frames within each chunk conditioned on both past context and future anchors. Importantly, these sparse anchors also provide geometric constraints, supported by bidirectional epipolar geometry, to localize where corresponding content should appear in the intermediate frames. Experiments on four benchmarks demonstrate consistent improvements over strong baselines in long-horizon visual quality, geometric consistency, and multi-view coherence. These gains further translate into improved downstream planning performance under the same planners, highlighting the importance of drift-resistant, geometry-aware prediction for reliable navigation world models.

Poster
Drift-React: One-step Generation of Reaction Pathways via SE(3) Drifting Fields

Rémi Schlama, Philippe Schwaller

Expand abstract

Mapping reaction pathways and transition states (TS) is fundamental to chemistry but computationally expensive at scale. The minimum energy pathway (MEP) dictates reaction rates and mechanisms, yet recovering it via electronic-structure methods requires thousands of costly force evaluations. Recent generative models accelerate TS identification but require slow iterative inference and only predict isolated saddle-point snapshots, missing the continuous reaction trajectory. We introduce Drift-React, an $\mathrm{SE}(3)$-equivariant generative framework that predicts complete reaction pathways in a single forward pass from only reactant and product geometries. By shifting distribution evolution to training via a Sinkhorn-weighted drifting field, Drift-React eliminates both the iterative force evaluations of NEB-style methods and the sequential ODE/SDE integration of diffusion and flow matching models. Evaluated on the Transition1x and Halo8 datasets, our one-step model generates physically consistent MEPs that accurately capture energetic bottlenecks and enable arbitrary-resolution sampling along the reaction coordinate. For isolated TS prediction, Drift-React matches the sub-Ångström accuracy of state-of-the-art iterative models while delivering orders-of-magnitude acceleration, clearing a major computational bottleneck for large-scale reaction network exploration.

Poster
Drift-Aware Multimodal User Representation Learning via Multi-Scale Temporal Modeling and Sparse Mixture-of-Experts

Ziqing Qian, Haohang Chen, Shengqi Dang, Yuhan Xiong, Canyu Shen, Jiaying Lei, Nan Cao

Expand abstract

Understanding user preferences from noisy and temporally evolving social media behaviors is fundamentally challenging due to interest drift, where user preferences shift across time and exhibit both multi-scale temporal patterns and diverse co-existing interests. To address this, we propose DUMoE, a unified framework for drift-aware multimodal user representation learning. Our model consists of (i) a temporal dynamics-aware backbone that captures and integrates static profiles, short-term behavioral signals, and long-term dependencies into a coherent representation, and (ii) a sparse mixture-of-experts (MoE) interest adapter that disentangles multiple latent interests via expert specialization and adaptive routing. Each expert models a distinct interest subspace, while a gating network dynamically selects and aggregates a sparse subset of relevant experts for each user. To enable stable and effective optimization, we further introduce a three-stage training strategy that decouples backbone learning, expert specialization, and gating optimization. Extensive experiments on real-world social media datasets show that DUMoE consistently outperforms state-of-the-art methods on both user interest prediction and interaction prediction tasks.

Poster
Drift Q-Learning

Achraf Anas El Houssaini, Mohamad Hosein Danesh, Amin Abyaneh, Scott Fujimoto, Hsiu-Chin Lin, David Meger

Expand abstract

Offline reinforcement learning requires improving a policy from fixed data while avoiding out-of-distribution actions with unreliable value estimates. Diffusion and flow policies handle this trade-off by modeling the behavior distribution to regularize the RL objective, but they require iterative denoising, solver integrations, and in more efficient variants, distillation or other approximations at inference. We propose DriftQL, which combines a drift-based behavioral regularizer with critic-driven policy improvement. The value signal biases the policy toward high-value regions of the data support, while attraction and repulsion together keep generated actions near the data and prevent collapse onto a single mode. DriftQL is implemented as a single network with a unified training objective and generates actions in a single forward pass. On D4RL and OGBench, DriftQL consistently outperforms diffusion and flow methods, advancing the state of the art. Under degraded data quality, where the baselines visibly struggle, DriftQL remains close to its clean-data performance, positioning it as a promising alternative to diffusion and flow-based methods while maintaining the simplicity and efficiency of deterministic approaches.

Poster
Drag as Evidence: Motion-Grounded Latent Recomposition for Drag-Based Editing

Xinyu Pu, Hongsong Wang, Jie Gui, Pan Zhou

Expand abstract

Modern image editors excel at semantic manipulation and visual synthesis, yet remain limited in precise spatial control, motivating the development of drag-based editing. However, existing drag-based methods often struggle to balance drag accuracy with natural, plausible, and intent-aligned generation. We propose MoRe-Drag, a motion-grounded drag-based editing method. Our key insight is to treat pixel-space warping as coarse motion evidence, and to inject this evidence into the generative sampling trajectory. Specifically, MoRe-Drag performs region-aware latent recomposition over refinement, inpainting, and anchor regions, coupled with stage-adaptive conditioning that progressively shifts from motion-grounded structure formation to semantic refinement. We further support an instruction-free interface by adapting the MLLM-based text encoder for drag-aware instruction inference. Experiments on DragBench-SR and DragBench-DR show that MoRe-Drag substantially improves drag precision over strong base editors and achieves superior drag accuracy among SOTA drag-based methods, while delivering strong semantic consistency and visually realistic results. Code and dataset will be publicly released.

Poster
DrPO: Drifting Preference Optimization for One-Step Generative Models

Zhou Jiang, Yandong Wen, Zhen Liu

Expand abstract

While one-step text-to-image generators, a fast and deployment-friendly class of visual generative models, can produce samples with a single forward pass, existing preference-finetuning methods fail to simultaneously achieve efficient adaptation, reward-model agnosticism, and diverse generation. In this work, we propose Drifting Preference Optimization (DrPO), an online preference-finetuning method for one-step generators. Inspired by the recent work of Drifting Models, DrPO first uses the target reward model to construct an on-policy dipole reward model from positive and negative samples. This dipole model induces a preference drifting field in latent feature space, whose gradients are then used to optimize the generator without backpropagating through the target reward model. Empirically, we evaluate DrPO across multiple reward models and SD-Turbo and SDXL-Turbo one-step generators, including HPSv3 and GenEval as representative alignment benchmarks, and validate that DrPO can efficiently and effectively finetune one-step generative models. Preliminary experiments further suggest that this sample-based gradient synthesis can be extended to offline preference finetuning.

Poster
Dr. MAS: Stable Reinforcement Learning for Multi-Agent LLM Systems

Lang Feng, Longtao Zheng, Shuo He, Fuxiang Zhang, Bo An

Expand abstract

Multi-agent LLM systems improve reasoning and tool use via role specialization, yet reinforcement learning (RL) post-training for such systems remains unstable and underexplored. We theoretically pinpoint a key source of instability when extending group-based RL to cooperative multi-agent LLM systems: under GRPO-style optimization, a global normalization baseline can mismatch heterogeneous agents' reward distributions, inducing gradient-norm instability. Based on this finding, we propose Dr. MAS, a simple and stable RL recipe for cooperative multi-agent LLM systems. Algorithmically, Dr. MAS normalizes advantages per agent using each agent's own reward statistics, which calibrates gradient scales and dramatically stabilizes the training. Systemically, Dr. MAS provides an end-to-end multi-agent RL framework with scalable orchestration, flexible per-agent LLM serving and optimization, and shared resource scheduling of actor backends. Unlike single-actor frameworks such as veRL, Dr. MAS natively supports multiple heterogeneous LLMs over a unified GPU pool. Lifecycle-aware backend management and dynamic dispatch release inactive models' GPU memory and schedule active agents on demand, enabling hardware-efficient co-training. We evaluate Dr. MAS on multi-agent math reasoning and multi-turn search with Qwen2.5 and Qwen3. Dr. MAS achieves clear gains over vanilla GRPO (e.g., +5.6\% avg@16 and +4.6\% pass@16 on math, and +15.2\% avg@16 and +13.1\% pass@16 on search) while largely eliminating gradient spikes. It also remains highly effective under heterogeneous agent-model assignments while improving efficiency. Code is available at https://anonymous.4open.science/r/DrMAS-0787.

Poster