ICML 2026 Papers

International Conference on Machine Learning

Decision breakdown

6,341papers

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

  • Spotlight
    5368.5%
  • Poster
    5,80591.5%

Browse all papers

Use What You Know: Causal Foundation Models with Partial Graphs

Arik Reuter, Anish Dhir, Cristiana Diaconu, Jake Robertson, Ole Ossen, Frank Hutter, Adrian Weller, Mark van der Wilk, Bernhard Schölkopf

Expand abstract

Estimating causal quantities traditionally relies on bespoke estimators tailored to specific assumptions. Recently proposed Causal Foundation Models (CFMs) promise a more unified approach by amortising causal discovery and inference in a single step. However, in their current state, they do not allow for the incorporation of any domain knowledge, which can lead to suboptimal predictions. We bridge this gap by introducing methods to condition CFMs on causal information, such as the causal graph or more readily available ancestral information. When access to complete causal graph information is too strict a requirement, our approach also effectively leverages partial causal information. We systematically evaluate conditioning strategies and find that injecting learnable biases into the attention mechanism is the most effective method to utilise full and partial causal information. Our experiments show that this conditioning allows a general-purpose CFM to match the performance of specialised models trained on specific causal structures. Overall, our approach addresses a central hurdle on the path towards all-in-one causal foundation models: the capability to answer causal queries in a data-driven manner while effectively leveraging any amount of domain expertise.

Poster
Hyperbolic Hierarchical Alignment for Video-Based Visible-Infrared Person Re-Identification

Shuang Li, Changjiang Kuang, Jiaxu Leng, Mingpi Tan, Zhanjie Wu, Shuanglin Yan, Xinbo Gao

Expand abstract

Video-based visible-infrared person re-identification (VVI-ReID) aims to learn robust video-level representations under modality discrepancy. However, existing methods typically rely on Euclidean geometry, which is suboptimal for modeling the complex temporal dynamics within visible and infrared tracklets, as it inevitably distorts the intrinsic hierarchical structure inherent in diverse temporal variations (e.g., occlusion, pose). In this paper, we propose Hyperbolic Hierarchical Alignment (HHA), which unifies spatio-temporal modeling and cross-modality alignment on the Poincaré ball. HHA employs a Hyperbolic Hierarchical Spatio-Temporal Aggregator (HHSA) to organize time-varying cues into low-distortion hierarchical representations via Hyperbolic Geometry Interaction (HGI) and Dual-Geometry Fusion (DGF). Furthermore, we introduce Geometry-Aware Modality Alignment (GMA), which integrates Hyperbolic Modality Alignment (HMA) to couple modality centroids for geometric consistency and Hyperbolic Prototype Alignment (HPA) to anchor both modalities to shared identity prototypes for robust discrimination. Experiments on HITSZ-VCM and BUPTCampus demonstrate state-of-the-art performance. The code will be available at https://github.com/Visuang/HHA.

Poster
Imitation Learning for Multi-turn LM Agents via On-policy Expert Corrections

Niklas Lauffer, Xiang Deng, Srivatsa R Kundurthy, Brad Kenstler, Jeff Da

Expand abstract

A popular paradigm for training LM agents relies on imitation learning, fine-tuning on expert trajectories. However, we show that the off-policy nature of imitation learning for multi-turn LM agents suffers from the fundamental limitation known as covariate shift: as the student policy's behavior diverges from the expert's, it encounters states not present in the training data, reducing the effectiveness of fine-tuning. Taking inspiration from the classic DAgger algorithm, we propose a novel data generation methodology for addressing covariate shift for multi-turn LLM training. We introduce on-policy expert corrections (OECs), partially on-policy data generated by starting rollouts with a student model and then switching to an expert model part way through the trajectory. We explore the effectiveness of our data generation technique in the domain of software engineering (SWE) tasks, a multi-turn setting where LLM agents must interact with a development environment to fix software bugs. Our experiments compare OEC data against various other on-policy and imitation learning approaches on SWE agent problems and train models using a common rejection sampling (i.e., using environment reward) combined with supervised fine-tuning technique. Experiments find that OEC trajectories show a relative 14% and 13% improvement over traditional imitation learning in the 7b and 32b setting, respectively, on SWE-bench verified. Our results demonstrate the need for combining expert demonstrations with on-policy data for effective multi-turn LM agent training.

Poster
FAFO: Lossy KV Cache Compression for Lossless Inference Acceleration via Draftless Fumble Decoding

Hoang Anh Duy Le, Shaochen Zhong, Yifan Lu, Yingtong Dou, Jiayi Yuan, Yu-Neng Chuang, Xiran Fan, Guanchu Wang, Yuzhong Chen, Xia Hu

Expand abstract

*Lossy KV cache compression* is a well-explored subfield of machine learning efficiency, with improved latency being one of its major gains. However, lossy compression techniques can fumble from time to time, exhibiting various, and often catastrophic, failure patterns that are not only difficult to resolve but sometimes even hard to identify, making direct deployment of models with compressed KV cache a risky endeavor. In this work, we explore a way to preserve *lossless* generation quality while still benefiting from the acceleration provided by KV cache compression. Specifically, we draw inspiration from the *n-gram candidate pool decoding* paradigm where we purposely allow the model to Fumble Around with compressed KV cache to generate multiple lossy "n-gram guesses", while in parallel Find Out via lossless verification in the same forward pass. From a conceptual standpoint, our proposed framework is compatible with all typical static or dynamic KV cache compression methods from the token dropping realm, thus opening up a new avenue for the stagnant n-gram decoding paradigm. Practically, we show that this framework presents many useful traits that similar draftless baselines (e.g., Self-Speculative Decoding) cannot achieve, such as requiring only one set of KV cache and being far less sensitive to model, task, and input-length scenarios. Our comprehensive empirical results show FAFO provides 1.20-2.71× latency speedup over the original model, while consistently outperforming other lossless + draftless solutions.

Poster
Mosaic: Unlocking Over 30$\times$ Context Length for Diffusion LLMs Inference via Global Memory Planning and Dynamic Peak Taming

Liang Zheng, Bowen Shi, Yitao Hu, Jiawei Zhang, Ruofan Li, Guotao Yang, Zhixin Zhao, Zhengchao Wang, Sheng Chen, Wenxin Li, Dezhi Ran, Tao Xie, Keqiu Li

Expand abstract

Diffusion-based large language models (dLLMs) have emerged as a promising alternative to autoregressive models, leveraging simultaneous denoising to enable global planning and iterative refinement. These properties make dLLMs attractive for long-context generation. However, deploying dLLMs faces a prohibitive memory barrier, as existing inference systems are inefficient for the diffusion paradigm. We observe that current inference systems are misaligned with dLLMs. Unlike autoregressive models, whose memory footprint is dominated by the KV-Cache, dLLMs are bottlenecked by transient activations rematerialized per step. Moreover, generic memory reuse mechanisms lack the global visibility to handle dynamic memory peaks of dLLMs, which alternate between logits and feed-forward networks. To address these challenges, we present Mosaic, a memory-efficient inference system that shifts dLLM execution from local, static management to a global, dynamic paradigm. Mosaic integrates (i) a mask-only logits kernel eliminating redundant activation materialization, (ii) a lazy chunking optimizer using online heuristics to tame dynamic memory peaks, and (iii) a global memory manager leveraging virtual addressing to mitigate memory fragmentation. Evaluations show that Mosaic reduces the memory peak-to-average ratio by 2.71$\times$ on average and increases the maximum inference sequence length on identical hardware by 15.30--32.34$\times$. Crucially, Mosaic is training-free and preserves exact model outputs, while reducing end-to-end latency by 2.5\%--55.4\%. Our code is publicly available at https://github.com/flashserve/Mosaic.

Poster
Reasoning to Edit: Hypothetical Instruction-Based Image Editing with Visual Reasoning

Qingdong He, Xueqin Chen, Chaoyi Wang, Yanjie Pan, Xiaobin Hu, Zhenye Gan, Chengjie Wang, Xiangtai Li, Jiangning Zhang, Yabiao Wang

Expand abstract

Instruction-based image editing (IIE) has advanced rapidly with the success of diffusion models. However, existing efforts primarily focus on simple and explicit instructions to execute editing operations such as adding, deleting, moving, or swapping objects. They struggle to handle more complex implicit hypothetical instructions that require deeper reasoning to infer plausible visual changes and user intent. Additionally, current datasets provide limited support for training and evaluating reasoning-aware editing capabilities. Architecturally, these methods also lack mechanisms for fine-grained detail extraction that support such reasoning. To address these limitations, we propose Reason50K, a large-scale dataset specifically curated for training and evaluating hypothetical instruction–reasoning image editing, along with ReasonBrain, a novel framework designed to reason over and execute implicit hypothetical instructions across diverse scenarios. Reason50K includes over 50K samples spanning four key reasoning scenarios: Physical, Temporal, Causal, and Story reasoning. ReasonBrain leverages Multimodal Large Language Models (MLLMs) for editing guidance generation and a diffusion model for image synthesis, incorporating a Fine-grained Reasoning Cue Extraction (FRCE) module to capture detailed visual and textual semantics essential for supporting instruction reasoning. To mitigate the semantic loss, we further introduce a Cross-Modal Enhancer (CME) that enables rich interactions between the fine-grained cues and MLLM-derived features. Extensive experiments demonstrate that ReasonBrain consistently outperforms state-of-the-art baselines on reasoning scenarios while exhibiting strong zero-shot generalization to conventional IIE tasks. Our dataset and code will be released publicly.

Poster
TexEditor: Structure-Preserving Text-Driven Texture Editing

Bo Zhao, Yihang Liu, Chenfeng Zhang, Huan Yang, Kun Gai, Wei Ji

Expand abstract

Text-guided texture editing aims to modify object appearance while preserving the underlying geometric structure. However, our empirical analysis reveals that even SOTA editing models frequently struggle to maintain structural consistency during texture editing, despite the intended changes being purely appearance-related. Motivated by this observation, we jointly enhance structure preservation from both data and training perspectives, and build TexEditor, a dedicated texture editing model based on Qwen-Image-Edit-2509. Firstly, we construct BlenderTex, a high-quality SFT dataset generated with Blender, which provides strong structural priors for a cold start. Secondly, we introduce StructureNFT, a RL–based approach that integrates structure-preserving losses to transfer the structural priors learned during SFT to real-world scenes. Moreover, due to the limited realism and evaluation coverage of existing benchmarks, we introduce TexBench, a general-purpose real-world benchmark for text-guided texture editing. Extensive experiments on existing Blender-based texture benchmarks and our TexBench show that TexEditor consistently outperforms strong baselines such as Nano Banana Pro. In addition, we assess TexEditor on the general-purpose benchmark ImgEdit to validate its generalization.

Poster
Curated Synthetic Data Doesn’t Have to Collapse: A Theoretical Study of Generative Retraining with Pluralistic Preferences

Ali Falahati, Mohammad Mohammadi Amiri, Kate Larson, Lukasz Golab

Expand abstract

Recursive retraining of generative models poses a critical representation challenge: when synthetic outputs are curated based on a fixed reward signal, the model tends to collapse onto a narrow set of outputs that over-optimize that objective, causing diversity to vanish and failing to represent the full range of preferences. Prior work has suggested that such collapse is unavoidable without adding real data into the mix. In this paper, we revisit that conclusion from an alignment perspective and show that collapse can be mitigated through curation based on multiple reward functions. We formalize the dynamics of recursive training under heterogeneous preferences and prove that, under certain conditions, the model converges to a stable distribution that allocates probability mass across competing high-reward regions. The limiting distribution preserves diversity and provably satisfies a weighted Nash bargaining solution, offering a formal interpretation of value aggregation in synthetic retraining loops.

Poster
When Softmax Fails at the Top: Extreme‑Value Corrections for InfoNCE

Hasan Sabri Melihcan Erol, Suat Evren, Oktay Ozel, Alexander Morgan, Jongha Jon Ryu, Lizhong Zheng

Expand abstract

InfoNCE is the standard contrastive learning objective, but its softmax form is not only a computational convenience: it also encodes a statistical assumption about how the top-scoring example is selected. Using extreme value theory, we show that this assumption is often misaligned with the normalized embedding setting used in modern contrastive learning. Motivated by this mismatch, we propose \textsc{WEINCE}, a simple modification of InfoNCE that uses anchor-wise online batch statistics to blend the usual softmax logits with an endpoint shortfall correction, adding no trainable parameters. Across five vision benchmarks, \textsc{WEINCE} yields consistent improvements in frozen-feature evaluation. These results show that a more faithful statistical treatment of hard negatives can improve contrastive objectives.\footnote{Code: \href{https://github.com/hsme98/weince}{github.com/hsme98/weince}.}

Poster
Robust Federated Learning Against Adaptive Compression

Wenjing Yan, Xiangyu Zhong, Ying-Jun Angela Zhang

Expand abstract

This paper addresses the critical challenges of hyperparameter tuning and communication efficiency in federated learning (FL). Despite recent advancements in parameter-free FL algorithms such as PAdaMFed, significant communication overhead remains a major obstacle to their practical deployment. To tackle these challenges, we propose a novel communication-efficient parameter-free FL algorithm ParFreFL that halves the communication requirements of PAdaMFed while preserving its parameter-free property. Building on this foundation, we introduce a compressed variant, ComParFreFL, which unifies the momentum increment and error feedback into a single parameter, effectively handling biased compression while maintaining the minimal communication cost. Notably, ComParFreFL also operates independent of the compression ratio, representing the first instance of such robustness in the compressed FL literature to our knowledge. Theoretically, our methods are proven to handle arbitrary data heterogeneity, partial client participation, and achieve linear speedup with respect to both local updates and participating clients. Extensive empirical evaluations demonstrate that our approaches match or slightly surpass the performance of carefully tuned alternatives while significantly reducing communication overhead, making FL more accessible and deployable in dynamic, resource-constrained environments.

Poster
Guided Star-Shaped Masked Diffusion

Viacheslav Meshchaninov, Egor Shibaev, Artem Makoian, Ivan Klimov, Nikita Balagansky, Daniil Gavrilov, Aibek Alanov, Dmitry Vetrov

Expand abstract

The performance of pre-trained masked diffusion models is often constrained by their sampling procedure, which makes decisions irreversible and struggles in low-step generation regimes. We introduce a novel sampling algorithm that works with pre-trained models and, after a lightweight fine-tuning of a single layer, significantly improves sample quality and efficiency. Our method reformulates the generation process using a star-shaped paradigm, which inherently allows for error correction. To make this process effective, we augment it with a learnable remasking module that intelligently identifies and revises likely errors. This approach yields a substantial quality boost, particularly when using a small number of sampling steps. We extensively ablate key components of our approach and show its usability in different scenarios. In experiments on text, and code generation, our sampling algorithm outperforms or matches existing methods. Code is available at https://github.com/EgorShibaev/G-Star.

Poster
UniDrag: Unified Multi-Field Prediction and Robust Shape Optimization for Vehicle Aerodynamics

Ye Liu, Shouyi Liu, Ding Wang, Huiyu Yang, Ruizhe Deng, Qian Li, Yuxiao Hu, Jianghang gu, Yongzheng Liu, Quanshi Zhang, Shiyi Chen, Yuntian Chen

Expand abstract

High-fidelity vehicle aerodynamics analysis is bottlenecked by costly CFD simulations. Neural surrogates accelerate prediction but lack inverse design capabilities, while existing generative optimization methods suffer from unstable convergence and frequent engineering constraint violations. We present UniDrag, a unified framework bridging multi-field aerodynamic prediction with robust differentiable shape optimization. Given a vehicle geometry, UniDrag predicts surface pressure, volume flow fields, drag coefficient $C_d$, and a streamwise build-up profile localizing drag contributions. Our architecture combines enhanced Physics-Sliced Attention (ePSA), Gated Expert Routing, and Modality-Protected Learning to prevent negative transfer across output modalities. At deployment, the frozen surrogate enables gradient-based optimization via Free-Form Deformation with engineering constraints. We introduce Expectation-over-Transformation to prevent adversarial exploitation of surrogate fragility. We curate a large-scale dataset of 15,000 vehicle geometries spanning four body types with GPU-accelerated LBM simulations. On this benchmark, UniDrag achieves $C_d$ prediction $R^2$ of 0.937 (+7.6\% over baselines) and 13.7\% mean CFD-verified drag reduction with 100\% success rate and only 21.3~mm average displacement (0.46\% vehicle length).

Poster
Quantitative Estimation of Target Task Performance from Unsupervised Pretext Task in Semi/Self-Supervised Learning

Lin-Han Jia, Siyu Han, Wen-Chao Hu, Jie-Jing Shao, Wen-Da Wei, Zhi Zhou, Lan-Zhe Guo, Yu-Feng Li

Expand abstract

The effectiveness of unlabeled data in Semi/Self-Supervised Learning (SSL) depends on appropriate assumptions for specific scenarios, thereby enabling the selection of beneficial unsupervised pretext tasks. However, existing research has paid limited attention to assumptions in SSL, resulting in practical situations where the compatibility between the unsupervised pretext tasks and the target scenarios can only be assessed after training and validation. This paper centers on the assumptions underlying unsupervised pretext tasks and explores the feasibility of preemptively estimating the impact of unsupervised pretext tasks at low cost. Through rigorous derivation, we show that the impact of unsupervised pretext tasks on target performance depends on three factors: assumption learnability with respect to the model, assumption reliability with respect to the data, and assumption completeness with respect to the target. Building on this theory, we propose a low-cost estimation method that can quantitatively estimate the actual target performance. We build a benchmark of over one hundred pretext tasks and demonstrate that our estimated performance strongly correlates with the actual performance obtained through large-scale training and validation.

Poster
Backjump-on-Graph: Empowering Large Language Models with Reinforced Retrospective Exploration for Agentic Knowledge Graph Reasoning

Yunqi Zhang, Shiqi Yan, Zhenzhao Yuan, Wenrui Liang, Yangming Liu, Zhixiao Qi, Tianyi Zhang, Shijie Zhang, Wei-Qiang Zhang, Yongfeng Huang, Haixin Duan, Shuai Chen, Yubo Chen

Expand abstract

Grounding Large Language Models (LLMs) in Knowledge Graphs (KGs) has shown significant promise for complex Question Answering (QA) tasks. Since LLMs' limited context window cannot accommodate the sheer volume of large-scale KGs, existing work usually utilizes agents to reason on real-world KGs, which follows reasoning paths derived from the queries step by step. However, the mismatch between query-derived paths and the KG's structure, stemming from users' lack of schema knowledge, usually leads the agents into dead ends. To address this problem, in this paper, we propose Backjump-on-Graph (BoG), a novel framework that empowers LLMs to retrospectively explore alternative reasoning paths at dead ends. We first propose to formalize each reasoning step with four atomic operations to create a structural scaffold that allows LLMs to revert to historical status. Next, we fine-tune the LLM with synthetic data containing the above atomic operations to instill basic backjump abilities. Finally, we leverage reinforcement learning and propose a hybrid reward function, which penalizes redundant transitions and promotes correct answers, to optimize the timing and landing nodes of backjumping. Extensive experiments on several KGQA benchmark datasets demonstrate the effectiveness of our BoG method.

Poster
Spectral Imbalance Causes Forgetting in Low-Rank Continual Adaptation

Hao Gu, Mao-Lin Luo, Zi-Hao Zhou, Han-Chen Zhang, Min-Ling Zhang, Tong Wei

Expand abstract

Parameter-efficient continual learning aims to adapt pre-trained models to sequential tasks without forgetting previously acquired knowledge. Most existing approaches treat continual learning as avoiding interference with past updates, rather than considering what properties make the current task-specific update naturally preserve previously acquired knowledge. From a knowledge-decomposition perspective, we observe that low-rank adaptations exhibit highly imbalanced singular value spectra: a few dominant components absorb most of the adaptation energy, thereby (i) more likely to disrupt previously acquired knowledge and (ii) making the update more vulnerable to interference from subsequent tasks. To enable explicit balance among components, we decouple the *magnitude* of the task update from its *directional structure* and formulate it as a constrained optimization problem on a restricted Stiefel manifold. We address this problem using a projected first-order method compatible with standard deep-learning optimizers used in vision-language models. Our method mitigates both backward and forward forgetting, consistently outperforming continual learning baselines. Source code is available in supplementary material.

Poster
Torus Graphs for Large Scale Neural Phase Analysis

Jack Goffinet, Casey Hanks, David Carlson

Expand abstract

Oscillatory neural signals such as electroencephalography (EEG) and local field potentials (LFPs) show phase relationships that coordinate communication across brain regions. Modern recordings capture hundreds of channels across many frequency bins, yet standard phase analyses are restricted to only a few variables. The Torus Graph (TG) model, an exponential-family distribution over phases whose univariate and pairwise potentials generalize von Mises distributions, infers principled structure among oscillations but models only static, undirected dependencies and is limited to $\sim \! 100$ variables because its score matching inference scales as $\mathcal{O}(d^{6})$. We introduce a stochastic score matching procedure that reduces the per-iteration cost to $\mathcal{O}(d^{2})$, enabling inference on datasets with thousands of variables. This scalable foundation supports analyses of 1,860 frequency-phase features from multi-electrode LFPs and enables two extensions previously inaccessible to TGs or classical circular statistics: (i) a TG Hidden Markov Model capturing state-dependent phase-coupling changes (e.g., spindle-related states during sleep) and (ii) an autoregressive TG inferring directional interactions via transfer-entropy estimation. Applied to LFP recordings, these models reveal state-dependent phase-interaction patterns between wakefulness and NREM sleep. Together, they enable systematic, large-scale mapping of dynamic and directional phase relationships across brain and cognitive states.

Poster
The Flexibility Trap: Rethinking the Value of Arbitrary Order in Diffusion Language Models

Zanlin Ni, Shenzhi Wang, Yang Yue, Tianyu Yu, Weilin Zhao, Yeguo Hua, Tianyi Chen, Jun Song, Cheng Yu, Bo Zheng, Gao Huang

Expand abstract

Diffusion Large Language Models (dLLMs) break the rigid left-to-right constraint of traditional LLMs, enabling token generation in arbitrary orders. Intuitively, this flexibility implies a solution space that strictly supersets the fixed autoregressive trajectory, theoretically unlocking superior reasoning potential. However, in this paper, we find that for general reasoning tasks (e.g., mathematics and coding), arbitrary order generation may in fact limit the reasoning potential of dLLMs. We observe that dLLMs tend to exploit this order flexibility to bypass high-uncertainty tokens that are crucial for exploration, which can lead to a premature collapse of solution coverage. This observation motivates a rethink of RL approaches for dLLMs, where considerable complexities, such as handling combinatorial trajectories and intractable likelihoods, are often devoted to preserving this flexibility. We show that effective reasoning can be elicited by simply forgoing arbitrary order and applying standard Group Relative Policy Optimization (GRPO) instead. Our approach, JustGRPO, is minimalist yet surprisingly effective (e.g., 89.1\% accuracy on GSM8K) while fully retaining the parallel decoding ability of dLLMs. Code: https://github.com/LeapLabTHU/JustGRPO.

Spotlight
The Relative Instability of Model Comparison with Cross-validation

Alexandre Bayle, Lucas Janson, Lester Mackey

Expand abstract

Cross-validation (CV) is known to provide asymptotically exact tests and confidence intervals for model improvement but only when the model comparison is *relatively stable*. Surprisingly, we prove that even simple, individually stable models can generate relatively unstable comparisons, calling into question the validity of CV inference. Specifically, we show that the Lasso and its close cousin, soft-thresholding, generate relatively unstable comparisons and invalid CV inferences, even in the most favorable of learning settings and when both models are individually stable. These findings highlight the importance of verifying relative stability before deploying CV for model comparison.

Spotlight
ScaleSim: Serving Large-Scale Multi-Agent Simulation with Invocation Distance-Based Memory Management

Zaifeng Pan, Yipeng Shen, Zhengding Hu, Zhuang Wang, Aninda Manocha, Zheng Wang, Zhongkai Yu, Yue Guan, Yufei Ding

Expand abstract

LLM-based multi-agent simulations are increasingly adopted across application domains, but remain difficult to scale due to GPU memory pressure. Each agent maintains private GPU-resident states, including models, prefix caches, and adapters, which quickly exhaust device memory as the agent count grows. We identify two key properties of these workloads: sparse agent activation and an estimable agent invocation order. Based on an analysis of representative workload classes, we introduce invocation distance, a unified abstraction that estimates the relative order in which agents will issue future LLM requests. Leveraging this abstraction, we present ScaleSim, a memory-efficient LLM serving system for large-scale multi-agent simulations. ScaleSim enables proactive prefetching and priority-based eviction, supports diverse agent-specific memory through a modular interface, and achieves up to 1.74$\times$ speedup over SGLang on simulation benchmarks. ScaleSim's source code is available at https://github.com/PanZaifeng/KVFlow.

Poster
LoPhyDA: Low-Rank Tensor and Physics Gradient Guided Diffusion for Atmospheric Data Assimilation

Danyang Peng, Yang Chen, Yunlong Zhou, Xiaotong Yuan

Expand abstract

Data Assimilation (DA) aims to integrate observations with model forecasts to estimate the state of dynamical systems. Despite the widespread application of diffusion-based assimilation methods, they remain constrained by the high dimensionality of atmospheric states and the reliance on imperfect state-observation mappings. This leaves regions lacking observations spatially unconstrained, leading to severe error accumulation and physical inconsistency.. In this paper, we propose LoPhyDA, a diffusion assimilation algorithm dual-guided by low-rank tensor and physical gradients. By leveraging the low-rank property of meteorological field, we employ tensor completion to exploit spatial continuity and dynamic correlations, reconstructing a globally informative dense field from sparse observations to serve as a global prior. This approach mitigates the information deficit inherent in sparse settings. The framework further incorporates physical constraints into the iterative denoising process, utilizing Partial Differential Equation (PDE) residual gradients to rectify the generative trajectory in real-time. Experimental results demonstrate that LoPhyDA outperforms state-of-the-art generative assimilation models in global weather prediction. It achieves robust and physically consistent assimilation, significantly reducing error accumulation in regions lacking observations.

Poster