Incomplete Data, Complete Dynamics: A Diffusion Approach
Authors: Zihan Zhou
Chenguang Wang
Hongyi Ye
Yongtao Guan
Tianshu Yu
Venue: ICLR 2026 Poster (2026)
Keywords: diffusion models
missing data
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: A diffusion framework learns complete dynamics from incomplete data via strategic context-query partitioning, with theoretical guarantees and strong imputation results on PDE and climate benchmarks.

Abstract: Learning physical dynamics from data is a fundamental challenge in machine learning and scientific modeling. Real-world observational data are inherently incomplete and irregularly sampled, posing significant challenges for existing data-driven approaches. In this work, we propose a principled diffusion-based framework for learning physical systems from incomplete training samples. To this end, our method strategically partitions each such sample into observed context and unobserved query components through a carefully designed splitting strategy, then trains a conditional diffusion model to reconstruct the missing query portions given available contexts. This formulation enables accurate imputation across arbitrary observation patterns without requiring complete data supervision. Specifically, we provide theoretical analysis demonstrating that our diffusion training paradigm on incomplete data achieves asymptotic convergence to the true complete generative process under mild regularity conditions. Empirically, we show that our method significantly outperforms existing baselines on synthetic and real-world physical dynamics benchmarks, including fluid flows and weather systems, with particularly strong performance in limited and irregular observation regimes. These results demonstrate the effectiveness of our theoretically principled approach for learning and imputing partially observed dynamics.

Published as a conference paper at ICLR 2026
INCOMPLETE DATA, COMPLETE DYNAMICS: A DIFFUSION APPROACH
Zihan Zhou!: Chenguang Wang 1, Hongyi Ye', Yongtao Guan', Tianshu Yu!
School of Data Science. The Chinese University of Hong Kong, Shenzhen Shanghai Artificial Intelligence Laboratory
{z (zihanzhoul, chenguangwang, hongyive] cuhk. edu. (guanyongaao, yutianshu) cuhk edd....kecccc... edu.cn
ABSTRACT
Learning physical dynamics from data is a fundamental challenge in machine learning and scientific modeling. Real-world observational data are inherently in- complete and irregularly sampled, posing significant challenges for existing data- driven approaches. In this work, we propose a principled diffusion- based frame- work for learning physical systems from incomplete training samples. To this end, our method strategically partitions each such sample into observed context and unobserved query components through a carefully designed splitting strategy, then trains a conditional diffusion model to reconstruct the missing query portions given available contexts. This formulation enables accurate imputation across ar- bitrary observation patterns without requiring complete data supervision. Specif- ically, we provide theoretical analysis demonstrating that our diffusion training paradigm on incomplete data achieves asymptotic convergence to the true com- plete generative process under mild regularity conditions. Empirically, we show that our method significantly outperforms existing baselines synthetic and real- world physical dynamics benchmarks, including fluid flows and weather systems, with particularly strong performance in limited and irregular observation regimes. These results demonstrate the effectiveness of our theoretically principled ap- proach for learning and imputing partially observed dynamics.
INTRODUCTION
Learning physical dynamics from observational data represents a cornerstone challenge in machine learning and scientific computing, with applications spanning weather forecasting (Conti 2024 Zhang et al. 2025b), fluid dynamics (Wang et al. 2024 Brunton & Kutz, 2024), biological systems modeling (Qi al. 2024 Goshisht (024). and beyond. Classical physics- based approaches require explicit specification of governing equations and boundary conditions, while data-driven methods offer the promise of discovering hidden dynamics directly from observations (Luo et al. 2025 Meng et al., 2025). However, a fundamental bottleneck persists: real-world observational data are inher- ently incomplete, irregularly sampled, and subject to various forms of missing information, making it difficult for existing approaches to learn accurate representations of the underlying dynamics.
Inherent sparsity of physical measurements. Physical science data fundamentally differs from typical computer vision datasets. Unlike natural images, where complete pixel grids are the norm, real-world physical measurements are inherently sparse and incomplete. Sensor networks provide observations only at discrete spatial locations, satellite imagery suffers from cloud occlusion, and experimental measurements are constrained by instrumental limitations. This incompleteness not temporary inconvenience to be resolved through better data collection is an intrinsic character- istic of how we observe physical systems.
Structured observation patterns. Prior approaches to learning from incomplete physical data have largely adopted simplistic assumptions about observation patterns. Most existing methods as- sume pixel-level independent and identically distributed (i.i.d.) missing patterns, where each spatial
"corresponding author

---

Published as a conference paper at ICLR 2026
location has equal probability of being observed (Daras al. 2023 Dai et al. 2024 Simkus & Gutmann 2025). While some recent works have explored alternatiee missing patterns in their experimental evaluations, such row/column missing for tabular data, they still employ the same training strategies regardless of the observation structure (Ouyang et al. 2023). This one-size-fits- all approach fails to leverage the specific characteristics of different mask distributions. In reality, observation patterns exhibit strong spatial structure: weather stations capture measurements within their local coverage areas, creating contiguous blocks of observations; satellite instruments observe swaths determined by orbital paths; underwater sensor arrays monitor volumes dictated by acoustic propagation. These structured patterns fundamentally differ from random pixel dropout and demand specialized training strategies. Our work addresses this gap by developing context-query partition- ing strategies specifically tailored to the underlying mask distribution, ensuring effective learning across diverse observation patterns.
Lack of theoretical foundations. While recent works have proposed various heuristic approaches for handling missing data in generative modeling, they lack rigorous theoretical foundations. Exist- ing methods typically rely empirical design choices without providing convergence guarantees or understanding of learning dynamics (Ouyang et al., 2023; Daras et al., 2023 Dai et al. 2024] Simkus & Gutmann 2025). Moreover, some theoretically-motive
approaches suffer from pro- hibitive computational costs, requiring multiple complete model retraining cycles or complex im- portance weighting schemes that limit their applicability to low-dimensional toy problems (Chen)
et al. 2024b
Givens et al. 2025 Zhang et al. 2025a).
Our solution. To address these challenges, develop a theoretically principled diffusion-based framework that provides rigorous convergence guarantees while maintaining computational effi- ciency for high-dimensional physical dynamics problems. Our approach answers critical questions about whether diffusion models trained solely on incomplete data can recover complete data dis- tributions, how observation patterns affect diffusion training efficiency, and under what conditions successful reconstruction of unobserved regions is guaranteed. In summary, our contributions are:
We propose a novel conditional diffusion training paradigm that works directly with incomplete training samples, featuring a strategically designed context-query par- titioning scheme tailored for physical dynamics.
Methodical design:
Theoretical guarantee: We provide the first theoretical analysis proving that ieffusioo-bbseed training on incomplete data with our paradigm asymptotically recovers the true complete dy- namical process under mild regularity conditions.
We conduct comprehensive experiments on both synthetic and real-world datasets, demonstrating substantial improvements in imputation accuracy over competitive baselines, particularly in challenging sparse observation regimes.
Strong results:
2 PRELIMINARIES
In Appendix A we present a review of imputation methods and generative modeling approaches for missing data, which provides the broader context for our contributions. We also provide de- tailed introduction diffusion models in Appendix B covering both standard noise matching and the data matching formulation that our method primarily employs. In this section, we formally define the problem of learning physical dynamics from incomplete observations. We establish the mathematical framework and notation that will be used throughout the paper.
We formalize the problem as follows. Let X C Rd denote the space of complete data samples fol- lowing unknown distribution Paaaa (Zo). Binary masks M {0, 1) are drawn from distribution Pmask( (M) where indicates observed elements. We assume that masks are conditionally indepen- dent of the data given the observation process: Pmmus (M xxn) Pmmask ( (M). In practice, we have prior knowledge about the mask distribution Pmis(M) based on the data collection process (e.g.,
sensor placement patterns, measurement protocols).
For each training instance i, we have x aos (1 -M ④ an representing missing values. Crucially, training dataset
M 6)
x denoting partially observed data
(6) and _. unobs
2

---

Published as a conference paper at ICLR 2026
D ((0 M contains only partial observations, complete samples xo are available during training. This setting reflects realistic scenarios where complete ground truth is unavailable. The objective is to learn conditional generative model pe( Zobs: Mf that generates complete samples consistent with the observed elements, despite being trained solely on incomplete data.
METHOD
In this section, we present our approach for learning physical dynamics directly from incomplete observations using diffusion models. Our method addresses the fundamental challenge of training generative models when both training and test data are partially observed, without access com- plete ground truth during training. Our approach consists of three key components: (1) Denoising data matching on incomplete training data (See. 3.T): We formulate a theoretically grounded training loss and establish conditions under which the model learns meaningful conditional ex- pectations for all dimensions. (2) Strategic context-query partitioning (Sec. 3.2): We develop principled strategy for partitioning incomplete samples into context and query components, en- abling reconstruction of originally missing dimensions. (3) Ensemble sampling for complete data reconstruction (Sec. 3.3): We bridge the gap between training on context masks and inference full observations through ensemble averaging with theoretical convergence guarantees. This unified framework enables robust learning from incomplete observations while providing strong theoretical foundations for reconstructing complete data from partial observations.
3.1 DENOISING DATA MATCHING INCOMPLETE TRAINING DATA
The fundamental challenge in learning from incomplete data is ensuring that training on incomplete observations enables the model to recover the complete underlying data distribution, despite never having access to complete samples during training. To address this, we formulate training objective that learns a conditional generative model pe (æo Tobs, MM 11 that generates complete samples con- sistent with the observed elements, despite being trained solely on incomplete data. Our approach strategically partitions incomplete samples and deliberately withholds information during training through a theoretically grounded framework.
Our key insight is to reframe the learning problem through a hierarchical masking strategy. For each incomplete sample (Fabs M)) in our training dataset, treat the partially observed data 2obs as "complete" within the scope of available observations. We then sample context masks Mex S M represent "observable" portions and query masks Mgr C M represent "query" portions relative to Lobs- Given the noisy sample Lobs,t M @ (at + o,é) at time i in the diffusion process, we train the neural network Ve to predict the complete clean data To from the timestep t, masked noisey observations M @ Lebs,t, and context mask Max. with our training loss is formulated as:
(tt, Bobs, MAx, Mgry) = |Mar @ (xr Max C Bobs,t, Max) xooss)l
(I)
The key architectural choice is that the model receives only the context mask Mex and correspond- ing observed values, trying to provide the best estimate for the queried dimensions. A natural ques- tion arises:
Since we train exclusively incomplete data, how this training loss enable the model predict the originally missing portions of the data -regions that were never observed during training?
Through the following theoretical analysis, we provide key insights into how training on incomplete observations can still lead to models capable of reconstructing complete data distributions.
Theorem 1 (Optimal solution under context masking without query information). Let z be the optimal solution by minimizing the loss in equation Under the conditional independence of masks and data, we have the following results:
For clarity, we note that in the full diffusion imputation setting, the ideal model would output E|zo Ct; Tebsa M). conditioning on both the noisy state x; and observations Tare. Our single-step sampling approach (Sec. 3.3) simplifies this by using minimal noise (t = 8 = 0), effectively approximating 15(wo Rots, M) For cases requiring diversity generation, provide multi-step sampling procedure in Appendix E that combines both sources of information through weighted averaging.
3

---

Published as a conference paper at ICLR 2026
(i) Optimal solution:
The optimal solution is given by
E ((xo), Mea O Lob Mai
P((Mgn)) P P(Mpp))
Maa) 0
(ro (t, Mas 3 Robs,ts Man));
arbitrary value,
Max)
where indicates the i-th entry of the vector. Specially, given the context mask Ma if the union of all possible query mask Mary supports covers all spatial dimensions, we have
De (1, Max Bobs,t, Mar.) = [20 Max O Bobs,ts Mad
(ii) Gradient magnitude scaling: The expected squared gradient magnitude with respect the net- work output for dimension scales linearly with the query probability ps == P((Mory) = 1 M
E (
1000 [
27
= -4p,B ((zo); - (Iana)i)? (Mary)s =1]
(4)
(iil) Parameter update frequency: The frequency of non-zero parameter updates for dimension exactly ps:
P(dimension i contributes to parameter update) = P((Mpr)) = M Pi
(5)
The proof can be found in Appendix This theorem reveals critical insight: the model learns meaningful conditional expectations E [mo); Mas 3 Lobs,t, Max] for dimension i only when P((May); Mex) O. When this probability is zero, the model's output for dimension becomes arbitrary since it never receives gradient updates for that dimension. This theoretical find- ing directly drives the need for strategic contexi-query partitioning: given any context mask Mes. (without information of M), dimensions that were originally missing in the raw data (i.c., dimensions i where M, = 0), has a positive probability of being selected query point. In the following section, detail how to design the context-query mask sampling strategy to guarantee this requirement.
must ensure that every dimension outside the context, including
3.2 STRATEGIC CONTEXT- PARTITIONING
Building on our theoretical analysis, we now address the crucial question: how should we design the context-query partitioning strategy to ensure effective learning? Our approach is guided by princi- pled design framework that guarantees positive query probabilities for all observable dimensions.
Design principle. Based Theorem context-query partitioning:
establish the core design principle for effective
Principle (Principle of uniform query exposure). For effective learning from incomplete data, the contexi-query partitioning strategy must satisfy:
1. Non-zero query probability: For all unobserved dimensions i, iee, (Mer); O,
P((Mr)); = Max) 00
(6)
2. Uniform exposure: The query probabilities should be approximately uniform across all ob- served dimensions to achieve balanced learning:
P((Mr)) = Mex)) Re P((Mgry) = Mew) Vi,j: (Mex); = (Mew), =
To implement this principle, we can decompose the query probability using the law of total proba- bility over all possible observation masks:
P((May): = Max) = Σ P((Map)) = Max. M) P(M Max)
(8)
M
This decomposition reveals that the query probability depends on two factors: (1) the conditional
query sampling strategy given both context and observation masks P((Ma = 1 Max- M), and (2) the posterior distribution observation masks given the context P(M Max). Since Max C M by construction,
always find observation masks M such that P(M Max) >
4

---

Published as a conference paper at ICLR 2026
Training Data with Added Noise
Context and Query Selection
Model Input and Learning Process
Model Output When Optimized
Pixel-level
Context: model input
P((M g0
Max)
-9 No loss update
Prodict arbitrary values
Query: calculate loss p((
Me..)
Possible loss update
Predict the expectations
Block-wise
Known mask distribution: complete blocks missing from 9-block grid
possible masks given Mars
((Mur)); = 1 Max)
P ((Mary)
Messo M
P (M
MMss)
Figure 1: Impact of context-query partitioning strategies on learning effectiveness. Blue regions indicate context (model input), orange regions indicate query (loss calculation). Top: Problematic uniform sampling creates zero-query regions. Bottom: Effective block- structrrd sampling ensures balanced learning across all dimensions. See Fig. 10 for the resulting imputation failures.
However, P((M = MAx. M) is not guaranteed to positive for all dimensions Fortunately, strategically design the sampling mechanism for Mex to ensure that there indeed exist observation masks M where both terms are simultaneously positive, thereby guaranteeing P((May); = Max) > 0 for all observed dimensions. To illustrate how different context mask selection strategies affect the query probability P( P((Mar = Max), we first present a concrete example that demonstrates the critical impact of this design choice.
Illustrate example. Fig. demonstrates how different partitioning strategies affect learning effec- tiveness using block-structured observation patterns. Consider scenario where observation masks M randomly mask 2 out of 9 spatial blocks:
Problematic strategy (top): For each partially observed sample Dobs, context points are selected by uniform sampling across all observable dimensions, yielding Musi This strategy is prob- lematic because for such given Musi there exists only one observation mask M that contains (i.e., the specific M from which Lobs was derived). Consequently, for masked dimensions where M, = 0, we have P((Mpry; = M M) 0 because these dimensions are never available for query selection. The regions marked with rectangles represent such zero- probability areas, leading to incomplete learning.
Effective strategy (bottom): Context masks are sampled by selecting entire blocks according to the same block-structured pattern as the observation masks, yielding Mbblck that typically con- tains 4 complete blocks. Crucially, for given Mbblck there exist multiple possible observation masks M that contain (see two visualized possible masks in the figure). This multiplicity ensures that for any masked dimensions i, there always exists at least one possible observation mask M such that P(( = bbllokk M) 0, thereby guaranteeing positive query probabilities across all observable dimensions.
Implementation strategy. Generally speaking, we typically have knowledge about how the data becomes masked (e.g.. sensor placement patterns, measurement protocols), which provides us with either explicit estimates reasonable prior knowledge about Pmask (M). In practice, to satisfy Prin-
5

---

Published as a conference paper at ICLR 2026
ciple [ viable strategy is to design the context mask sampling mechanism Max based the observation mask distribution Pmask [M). Specifically, we sample Mex from M following the same structural pattern as Pmask M): for iid. pixel-level observations, we independently sample each observed pixel; for block-structured observations, we sample complete blocks from available blocks in M. This distribution-pressene be excluded from context (and included in query), guaranteeing P ((Mgry); = Max) > 0 for all Under this design paradigm, the model learns during training recover 2o from context obser- vations while leveraging the structural knowledge embedded in Pmask (M.. At inference time, the model can then utilize the complete observation Lobs along with the same distributional knowledge Pmask (M) to reconstruct the full data Zo.
strategy ensures that every observed dimension can potentially
Training algorithm and practical considera- tions. Building our theoretical foundation and design principles, we present our training al- gorithm in Alg. [ While ensuring P((Mar) Max) > is necessary, the choice of context- query ratio involves two critical trade-off's
Algorithm 1 Diffusion- training for miss- ing data imputation
Require: dataset D = { M }
Ensure: trained neural network parameters O
I: while not converged do for each batch (Iebs; M) do
Information gap trade-off (Theorem 2
When Mex contains few observed points relative to M, the large information gap in- creases approximation variance and slows down convergence.
3: sample ~ Uni.(0, 1), N (0,
Bobs, <- M ② (ar,2r + 0,E) sample M M May S M (Princ.
5:
6:
Φ f- xe(t, Max C Tobs Max)
Parameter update frequency trade-off (equa- tion 5): When Max contains too many ob- served points, query probabilities become small, leading to infrequent parameter up- dates for reconstructing missing informa- tion.
C+ || (Ø Eobs)
8: update 0 using gradient descent on L
end for end while
10:
These theoretical considerations suggest that moderate context ratios should achieve optimal per- formance by balancing both trade-offs. Our experimental results confirm this theoretical prediction, with detailed analysis provided in Appendix H.4 and Tab. 7
3.3 ENSEMBLE SAMPLING FOR COMPLETE DATA RECONSTRUCTION
Our trained model approximates the conditional expectation E [pp Max Do Max] given a ran- domly sampled context mask, rather than the desired full conditional expectation E|a0 Fabs,t> M] that conditions on the complete observation. To bridge this gap and enable complete data recon- struction, we leverage ensemble averaging across multiple context masks. This section presents our sampling procedures and their theoretical guarantees.
Single-step sampling. In many scientific applications, the observed data are sufficiently informa- tive to constrain the solutions to relatively concentrated region in the solution space (Alberti & Santacesaria, 2021). When the posterior distribution p(zo &ooss M) is highly concentrated with solutions clustered closely together, it can be well-approximated by narrow distribution centered at x", where x" represents the mean of the tightly clustered solutions consistent with the observations. In such cases, 1B(xg Tabs M) x x² provides good representative solution, reducing the need for extensive iterative denoising steps.
To leverage this property, we implement single-step sampling procedure. We apply minimal noise <
at timestep t where
Osgaaas + OBE,
<~N(0,,T)
The small noise level ensures M O xs Lobs. We then approximate the desired conditional expectation using ensemble averaging over K randomly sampled context masks:
K
x²
= E
E[zo
Eobs
M
~
Σ
xg
M o Dobs,6, M(()
(10)
K

---

Published as a conference paper at ICLR 2026
M conditionally independent given M. This approach enables direct reconstruction in single denoising step, making it particularly suitable for well-posed inverse problems where the observations strongly constrain the solution space.
where M
M M
We then establish the theoretical foundation that justifies our ensemble averaging approach. Let obs M @ Tobs,r, MM and ctx [Mxx C Robs,t, Max] denote the full observation and context observation respectively.
Theorem 2 (Ensemble approximation convergence). Let El obs) == E Ea Dobs.t M be the ground truth conditional expectation and [F( ctx] = E El M O Robs, Menl be the context-conditioned expectation. The expected squared error between these quantities is:
E [|
ctx] El obs]]| = E[Var|æp ctx]] - E[Var[ | obs]]
(11)
Consider practical model with oulput æo po,,,ctx) Elaro ctx] b((lx) + Emias
Esia (ctx), where b(ctx)
represents context-dependent deterministic bias and Exias (ctx) is random error with E| E[Epiar) = 0. Given the ensemble prediction in equation T0aa:
K
j.K
Do ctx!
(12)
K
the expected squared error between the ensemble prediction and ground truth is: ||áx - E2 obs]]|
|| [
ctx]]
FF obs] + [[b(ctx)] 1|?
information gop
model higs
(13)
Var [
etx]] + Var Var[b(ctx))
+ Var(enia)
data variance
model variance
As K -> 00, the ensemble converges to:
lim K
|| EEzo
obs]]l] = |||[]
ctx]]
-Baoo
obs] + B[b(ctx)]]1
(14)
This theorem demonstrates that ensemble averaging eliminates the variance terms, with the remain- ing error determined by the information gap between context and full observations, plus any sys- tematic model bias. The proof is provided in Appendix F2
Multi-step sampling. While single-step sampling suffices when observations nearly determine a unique solution, generating diverse imputed samples handling cases with significant uncertainty requires a multi-step sampling procedure. This approach follows the standard diffusion sampling process but replaces each denoising step with ensemble averaging over multiple context masks. Note that multi-step sampling involves repeated application of the model to generated content, which can lead to slight accumulation of errors compared to the single- step approach (Xu et al. 2023). The detailed multi-step sampling algorithm is provided in Appendix
4 EXPERIMENT
4.1 BASELINES
We compare against established baselines, including traditional imputation methods (Temporal Con- sistency (Huang et al.] 2016), Fast Marching (Telea) 2004), Navier-Stokes inpainting (Bertalmio al. 2001h) and recent diffusion-based approaches, MissDiff (Ouyang et al., (023). AmbientD- ilf (Daras et al. 2023). To ensure a fair comparison, we modified MissDiff to use data matching instead of its original noise matching approach, which improves its performance (see Appendix H.. for more detailed discussion). As result, all three diffusion-based approaches, MissDifff Ambi- entDiff, and our method, now employ the data matching paradigm. For baseline methods, we experimented with both single-step sampling and multi-step sampling strategies, and report the best performance results between these two approaches to ensure optimal baseline comparisons. We also exclude methods by Chen et al. (2024b) Givens et al.] (2025); Zhang et al. (2025a) due to their computational limitations: these approaches are designed for fow-dimensional data and incur prohibitively high training costs when scaled to high dimensions. See Appendix G for details.

---

Published as conference paper at ICLR 2026
T-0
T-16
T-33
T-49 T-66 T=82 T-99
Original Data
Imputed Data
Figure 3: Comparison of original and imputed data from the Navier-Stokes dataset (60% observed points). The upper row shows the original data, while the bottom row shows the results after data imputation. Each sample consists of 100 frames 64 64 resolution.
4.2 DATASETS
We conduct comprehensive experiments to validate our theoretical framework and demonstrate the effectiveness our approach across diverse scientific domains. Our evaluation encompasses both synthetic PDE datasets: Shallow Water (Klöwer et al. 2018), Advection (Klöwer et al. 2018), and Navier-Stokes Cao 2024). and real-world climate data (ERA5) (Hersbach et al. (020), under vary- ing levels of data sparsity, ranging from 80% as low 1% observed points. To simulate realistic scientific measurement scenarios, we construct datasets where each sample contains only a subset of spatial locations with known values, while the remaining locations are permanently unobserved. This reflects the fundamental challenge in scientific applications where complete ground truth data never available during training, distinguishing our setting from conventional imputation tasks that artificially mask complete observations.
Shallow Water and Advection Equations. We consider two fundamental geophysical PDE sys- tems: the shallow water equations governing fluid dynamics with rotation, and the linear advection equation describing scalar transport. Each dataset contains 5k training, 1k validation, and Ik test sam- ples with 32 x 32 spatial resolution and 50 temporal frames, generated with randomized physical param- eters and initial conditions.
(a) Visualization data imputation results on Shallow Water and Advection datascts. The train- ing set contains 30% observed data points. Each example shows complete sample reconstructed by the model based on partial observations. Ob- served entries are marked with rectangles, while missing entries are filled in by the model.
Navier-Stokes Equations. We use incompressible Navier-Stokes simulations of isotropic turbulence, featuring the characteristic Kolmogorov energy cas- cade. The dataset comprises 1,152 samples at 64 x 64 resolution with 100-frame sequences, generated using spectral or finite volume methods.
Imputed Data
Original Data
ERAS Reanalysis. For real-world evaluation, we utilize the ERA5 atmospheric reanalysis from ECMWF, incorporating nine essential meteorologi- cal variables. We process one year of hourly data at 103 120 spatial resolution, segmented into 3-hour windows with sparse observations.
We assess reconstruction quality using physically meaningful metrics: PDE residual errors for shallow water, forward propagation accuracy for advection, and direct MSE for Navier-Stokes and ERA5. De- tailed dataset specifications and evaluation protocols provided in Appendix H.1
(b) Visualization of data imputation results "total column water vapor" variable in the ERA5 dataset (20% observed data points).
Figure Data imputation visualization
8

---

Published as a conference paper at ICLR 2026
Table Performance comparison on physical dynamics imputation tasks, where masks are sampled pixel-wisely. Column headers indicate the percentage of spatial points observed in the dataset.
Navier-Stokes (x10-))
ERA5 (x10
Method
80%
60%
20%
20%
10%
1%
Temporal Consistency
1.341
2.709
5.709
0.967
1.179
9.735
Fast Marching
0.486
1.220
3,737
0.710
0.978
3,053
Navier-Stokes
0.263
0.656
2.989
0.600
0.942
3.074
MissDiff
0,251 = 0.02s
0.611
= 0,077
3.077
1.046
0,416 = 0,004
0.676 0.08s
1.653 + 0.296
AmbientDiff
0.238
0.017
0.538
0.004
2.043
0.089
0.256
1: 0.002
0.414
0031
1.234
0:437
Ours
0,223 1: 0.016
0.507 1 0,826
1.931 + 0.092
0.250 1 0,002
0,408 0,030
1.229 1 0.437
Table 2: Performance comparison on PDE imputation tasks, where masks are sampled block-wisely. Column headers indicate the fraction of observed blocks. The last two rows demonstrate our method using incorrect pixel-level versus correct block-win wwise context-query partitioning strategies.
Shallow Water
Advection
Navier-Stokes
Method
8/9
5/9
8/9
Temporal Consisteney
0.6974
2.5486
0,4758
1.0940
1.4287
Fast Marching
0,8454
3,3718
0,5042
1.4434
1.7391
Navier-Stokes
0,4753
1,7565
0,4418
1,3594
1.7274
MissDiff'
0,0285
0.0024
0,1166
0.0066
0,1202
0.047
0.1979
0.0229
1.4357
0.1132
AmbientDiff'
0.0217 - 0.0033
0.0925
9.0017
0.1077
0.1524
1.4954
Ours
pixel-level (incorrect)
block-wise (correct)
0.0215 +
0.0203
0,005)
0.0989
0.0865
9,007
0.0114
0.1171
0.1065
1
0.0009
0.1894
0.1407 + 0.0116
0.7592
1.4925
0.1986
4.3 EXPERIMENT RESULTS
Individual spatial points are randomly masked throughout the domain, simulating sparse sensor networks or measurement failures. We test observation rates from 80% down to 1,, challenging the model to reconstruct scattered missing points using local spatial cor- relations. Tab. | and Tab. [0] demonstrate that our method achieves superior performance in the majority of evaluation scenarios.
Pixel-level observation.
Block-wise observation. This more challenging setting masks entire contiguous spatial regions, reflecting realistic constraints such as sensor placement limitations, regional measurement failures, or structured occlusions in observational systems. For instance, in the 5/9 block configuration, only 5 out of 9 spatial blocks contain observations, while 4 complete blocks remain entirely unobserved. This requires the model to reconstruct entire spatial regions without any local observations, relying solely on distant context and learned physical priors. Our results, shown in Tab. 2 demonstrate that the proposed strategic context-query partitioning, which adapts to the observation pattern dur- ing training, is essential. When the partitioning strategy matches the observation structure (block masks), our method effectively learns to reconstruct complete fields. Conversely, mismatched strate- gies lead to degraded performance, validating our theoretical analysis.
Table 3: Cross- dinsssdistrittiin generalization on the Navier-Stoker dataset. Each column represents a model trained with a specific observation ratio, and each row represents the test observation ratio. Values indicate MSE between reconstructed and ground truth fields. The diagonal entries represent matched train-test distributions, while off-diagonal entries measure generalization under distribution shift. Models maintain reasonable performance when test-time observations are close to training conditions, but degrade gracefully when trained on fewer observation datasets.
Test Set Training Set
80%
60%
20%
80%
0.2229 + 0.0122
0.2362 + 0.0121
0.3363 + 0.0082
60%
0.4990 + 0.0260
0.5071 + 0.0257
0.6980 + 0.0188
20%
1.9315 + 0.0921

---

Published as a conference paper at ICLR 2026
Cross-distribution generalization. To evaluate the robustness of method under distribution shift between training and testing, we conduct experiments where the observation ratio differs be- tween training and inference. Specifically, we investigate whether model trained on data with certain observation density can generalize to test scenarios with different observation patterns. The key challenge lies in maintaining consistent model behavior when the available information at test time deviates from the training distribution. Our implementation addresses this through adaptive context mask sampling: during training with observation ratio Thaii (c.g., 80%), we sample context masks containing a fraction of the observed points (e.g.. 50%), resulting in the model receiving Turann X a of the total pixels as input (e.g.. 40%). At test time with different observation ratio Mest (e.g., 60%), we maintain the same effective input ratio by sampling
Tinis Xa of the available obser-
Yiza
40% = 6%). This strategy ensures the model operates within its learned vations as context (e.g., 60%
input distribution while adapting to varying observation densities. Tab. 3 presents results across different train-test observation ratio combinations, demonstrating that our method maintains reason- able performance even under significant distribution shifts, though performance naturally degrades when the test-time observation ratio substantially deviates from the training distribution.
4.4 ABLATION STUDY
We conduct comprehensive ablation studies to analyze the contribution of key components in our framework: (1) test-time gap introduced by replacing Max with M, which evaluates the distribution mismatch between training and inference masks and its effect on reconstruction accuracy: (2) the choice of backbone architecture, where we investigate whether the proposed training paradigm gen- eralizes across different neural architectures; (3) the context and query mask ratio selection guided by our theoretical analysis, examining the trade-off between information availability and parameter update frequency during training; and (4) the influence of the ensemble size K, which controls the variance reduction effect in our ensemble approximation during sampling. Detailed results and anal- ysis are provided in Appendix H.4 Through systematic ablation studies, we validate our key theoret- ical insights. The results demonstrate that our method consistently outperforms existing approaches while providing theoretical guarantees for convergence to the desired conditional expectations.
CONCLUSION
We presented principled framework for learning physical dynamics from incomplete observations using diffusion models trained directly on partial data through strategic context-query partitioning. Our approach enables diffusion models learn consistent conditional expectations without requir- ing access to fully observed training samples, making it particularly suitable for realistic scientific datasets where missing data is unavoidable. Our theoretical analysis proves that training on incom- plete data recovers the complete distribution with convergence guarantees, validated empirically with substantial improvements over baselines on synthetic PDEs and ERA5 climate data, especially in sparse regimes (1-20% coverage). The method's ss effectiveness across diverse physical systems demonstrates practical applicability for real-world scenarios where complete observations are inher- ently unavailable. More broadly, our results highlight the potential of diffusion-based generative models as a powerful tool for modeling partially observed physical systems and open promising directions for future research on integrating generative modeling with scientific data analysis.
REPRODUCIBILITY STATEMENT
To ensure the reproducibility of our work, all experimental code and datasets are publicly available at https:://ithhn com/ LOGO-CUHKSZ / www.monnoneeeeee/emmmmmmcomm.eee data- complete- dameee-ddmmmmmmmmmm The complete implementation, including model architectures, training procedures, and evaluation scripts, is provided in the repository with detailed documentation. All datasets used in our experi- ments are released, with the exception of the ERA5 dataset, which requires individual registration and download from the official ECMWF website due to licensing restrictions.
ACKNOWLEDGMENTSS
This work was supported by the National Science and Technology Major Project under Grant 20227.001 16408.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Michael Albergo, Mark Goldstein, Nicholas M Boffi, Rajesh Ranganath, and Eric Vanden-Eijnden. Stochastic interpolants with data-dependent couplings. arXiv preprint arXiv:2310.03725, 2023.
Giovanni S Alberti and Matteo Santacesaria. Infinite dimensional compressed sensing from anisotropic measurements and applications to inverse problems in pde. Applied and Computa- tional Harmonic Analysis, 50:105-146, 2021.
Marcelo Bertalmio, Andrea L Bertozzi, and Guillermo Sapiro. Navier-stokes, fluid dynamics, and image and video inpainting. In Proceedings of the 2001 IEEE Computer Society Conference on Computer Vision and Pattern Recognition. CVPR 2001, volume 1. pp. I-I. IEEE, 2001.
Christopher M Bishop and Nasser M Nasrabadi. Pattern recognition and machine learning, vol- ume 4. Springer, 2006.
Steven L Brunton and J Nathan Kutz. Promising directions of machine learning for partial differen- tial equations. Nature Computational Science, 4(7):483-494, 2024.
Shuhao Cao. Navier-stokes dataset of isotropic turbulence in a periodic box, 2024. URL https: Funded by
//buggingfaat agw.ae..e........cooo http:////...cccccc.c.co.c/tttttt
National Science Foundation: NSF award DMS-2309778.
Hanyang Chen, Yang Jiang, Shengnan Guo, Xiaowei Mao, Youfang Lin, and Huaiyu Wan. Dif- flight: partial rewards conditioned diffusion model for traffic signal control with missing data. Advances in Neural Information Processing Systems, 37:123333-12337 2024a.
Zhichao Chen, Haoxuan Li, Fangyikang Wang, Odin Zhang, Hu Xu, Xiaoyu Jiang, Zhihuan Song, and Hao Wang. Rethinking the diffusion models for missing data imputation: A gradient flow perspective. Advances Neural Information Processing Systems, 37:1 12050-1121000 2024b.
Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Mare Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. arXiv preprint arXiv:2209. 2022.
Silvia Conti. Artificial intelligence for weather forecasting. Nature Reviews Electrical Engineering, 1(1):-88, 2024.
Ciprian Corneanu, Raghudeep Gadde, and Aleix M Martinez. Latentpaint: Image inpainting in latent space with diffusion models. In Proceedings of the IEEE/CVF winter conference on applications of computer vision, pp. 4334-4343, 2024.
Zongyu Dai, Emily Getzen, and Qi Long. Sadi: Similarity-awarr diffusion model-based imputa- tion for incomplete temporal ehr data. In International Conference Artificial Intelligence and Statistics, pp. 4195-4203, PMLR, 2024.
Giannis Daras, Kulin Shah, Yuval Dagan, Aravind Gollakota, Alex Dimakis, and Adam Klivans. Ambient difffusion: Learning clean distributions from corrupted data. Advances in Neural Infor- mation Processing Systems, 36:288-313, 2023.
Yilun Du, Katie Collins, Josh Tenenbaum, and Vincent Sitzmann. Learning signal-agnostic mani- folds of neural fields. Advances in Neural Information Processing Systems, 34:8320-833 -3311, 2021.
Yifan Duan, Jian Zhao, Junyuan Mao, Hao Wu, Jingyu Xu, Caoyuan Ma, Kai Wang, Kun Wang, Xuelong Li, et al. Causal deciphering and inpainting in spatio-temporal dynamics via diffusion model. Advances in Neural Information Processing Systems, 37:107604-107632, 2024.
Emilien Dupont, Hyunjik Kim, SM Eslami, Danilo Rezende, and Dan Rosenbaum. From data to functa: Your data point is a function and you can treat like one. arXiv preprint 22X1::2001.12204, 2022.
Bradley Efron. Tweedie's formula and selection bias. Journal of the American Statistical Associa- tion, 106(496):1002-6614 2011.
Christoph Feichtenhofer, Yanghao Li, Kaiming He, et al. Masked autoencoders as spatiotemporal learners. Advances in neural information processing systems, 15:35946-35955 2022.
11

---

Published as a conference paper at ICLR 2026
Josh Givens, Song Liu, and Henry WJ Reeve. Score matching with missing data. arXiv preprint arXiv:2506.00557. 2025.
Manoj Kumar Goshisht. Machine learning and deep learning in synthetic biology: Key architectures, applications, and challenges. ACS omega, 99))9922 -9945, 2024.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked au- toencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 16000-16009, 2022.
Hans Hersbach, Bill Bell, Paul Berrisford, Shoji Hirahara, András Horányi, Joaquín Muñoz-Sabater Julien Nicolas, Carole Peubey, Raluca Radu, Dinand Schepers, et al. The era5 global reanalysis. Quarterly journal of the royal meteorological society, 1469730):1999-204, 2020.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in
neural information processing systems, 2020.
Jia-Bin Huang, Sing Bing Kang, Narendra Ahuja, and Johannes Kopf. Temporally coherent com- pletion of dynamic video. ACM Transactions on Graphics (ToG), 15
Jiahe Huang, Guandao Yang, Zichen Wang, and Jeong Joon Park. Diffusionpde: Generative pde- solving under partial observation. arXiv preprint 24t0:2200.17766. 2024.
Niels Bruun Ipsen, BeereeAAexandrr Mattei, and Jes Frellsen. not-miwae: Deep generative mod- elling with missing not at random data. arXiv preprint arXiv:2006.12877, 2020.
Tero Karras, Mika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyz- ing and improving the training dynamics of diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 24174 2024.
Kwanyoung Kim and Jong Chul Ye. Noise2score: tweedie's approach self-supervised image denoising without clean images. Advances in Neural Information Processing Systems, 34:864- 874, 2021.
M. Klöwer, M. F. Jansen, M. Claus, R. Greatbatch, and Thomsen. Energy budget-based backscatter in shallow water model of a double gyre basin. Ocean Modelling, 132, 2018. doi: 10.1016/)..cemoo.2200 2018..0........ 0.09.066.
Jean-Marie Lemercier, Julius Richter, Simon Welker, Eloi Moliner, Vesa Välimäki, and Timo Gerk- mann. Diffusion models for audio restoration: review [special issue on model-based and data- driven audio signal processing]. IEEE Signal Processing Magazine, 41(6):72-84, 2025.
Steven Cheng-Xine Xian Li and Benjamin Marlin. Learning from irreuulatly-sumpled time series: A missing data perspective. In International Conference on Machine Learning, pp. 5937-5946. PMLR, 2020.
Steven Cheng-Xian Li, Bo Jiang, and Benjamin Marlin. Misgan: Learning from incomplete data with generative adversarial networks. arXiv preprint arXiv: 2019.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential cquations. arXiv preprint arXiv:2010.08895 2020.
Guan-Homg Liu, Arash Vahdat, De-An Huang, Evangelos A Theodorou, Weili Nie, and Anima Anandkumar. Rss:: Image-to-image schrödinger bridge. arXiv preprint mrxXi::3302.0872.. 2023.
Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool. Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 11461-11471, 2022.
Kuang Luo, Jingshang Zhao, Yingping Wang, Jiayao Li, Junjie Wen, Jiong Liang, Henry Sockmadji, and Shaolin Liao. Physics-informed neural networks for pde problems: a comprehensive review. Artificial Intelligence Review, 58 10):1-43, 2025.
12

---

Published as a conference paper at ICLR 2026
Chao Ma, Sebastian Tschiatschek, Richard Turner, José Miguel Heráandez-Lobat, and Cheng Zhang. Vaem: a deep generative model for heterogeneous mixed type data. Advances in Neural Information Processing Systems, 33:1 1237-1124711 2020.
Chuizheng Meng, Sam Griesemer, Defu Cao, Sungyong Seo, and Yan Liu, When physics meets machine learning: A survey of physics-informed machine learning. Machine Learning for Com- putational Science and Engineering, 1(1):00, 2025.
Yidong Ouyang, Liyan Xie, Chongxuan Li, and Guang Cheng. Missdiff: Training diffusion models on tabular data with missing values. arXiv preprint arXiv:2307.00400 00467, 2023.
Xin Qi, Yuanchun Zhao, Zhuang Qi. Siyu Hou, and Jiajia Chen. Machine learning empowering drug discovery: Applications, opportunities and challenges. Molecules, 29(4):903, 2024.
Chitwan Saharia, William Chan, Huiwen Chang, Chris Lee, Jonathan Ho, Tim Salimans, David Fleet, and Mohammad Norouzi. Palette: Image-to-image diffusion models. In ACM SIGGRAPH 2022 conference proceedings, pp. 1-10, 2022.
Vaidotas Simkus and Michael U Gutmann. Cfmi: Flow matching for missing data imputation. arXiv preprint arXiv:2506. 250600928, 2025.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Yang Song. Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 2020.
Alexandru Telea An image inpainting technique based on the fast marching method. Journal of graphics tools, (1):23-34, 2004.
Haixin Wang, Yadi Cao. Zijie Huang, Yuxuan Liu, Peiyan Hu, Xiao Luo. Zezheng Song, Wanjia Zhao, Jilin Liu, Jinan Sun, et al. Recent advances on machine learning for computational fluid dynamics: A survey. arXiv preprint arXiv:2408. /2171, 2024.
Yilun Xu, Mingyang Deng, Xiang Cheng. Yonglong Tian, Ziming Liu, and Tommi Jaakkola. Restart sampling for improving generative processes. Advances in Neural Information Processing Sys- tems, (6:76606--68336 2023.
Kaiwen Xue, Yuhao Zhou, Shen Nie, Xu Min, Xiaolu Zhang, Jun Zhou, and Chongxuan Li. Unify- ing bayesian flow networks and diffusion models through stochastic differential equations. arXiv preprint arXiv:2404.11 15766, 2024.
Conghan Yue, Zhengwei Peng, Junlong Ma, Shiyan Du, Pengxu Wei, and Dongyu Zhang, Im- age restoration through generalized ornstein- umstein-ublenbeck bridge. arXiv preprint arXiv:2312.10299, 2023.
Hengrui Zhang, Liancheng Fang, Qitian Wu, and Philip Yu. Diffputer: Empowering diffusion models for missing data imputation. In The Thirteenth International Conference Learning Representations, 2025a.
Huijun Zhang, Yaxin Liu, Chongyu Zhang, and Ningyun Li. Machine learning methods for weather forecasting: A survey, Atmosphere, 16(1):82, 2025b.
Kaiwen Zheng, Cheng Lu, Jianfei Chen, and Jun Zhu. Improved techniques for maximum likelihood estimation for diffusion odes. In International Conference on Machine Learning, pp. 42363- 42389. PMLR, 2023.
Zihan Zhou, Xiaoxue Wang, and Tianshu Yu. Generating physical dynamics under priors. arXiv preprint arXiv:2409.007333 2024.
Peiye Zhuang, Samira Abnar, Jiatao Gu, Alex Schwing, Joshua M Susskind, and Miguel Angel Bautista. Diffusion probabilistic fields. In The Eleventh International Conference on Learning Representations, 2023.
13

---

Published as a conference paper at ICLR 2026
A RELATED WORK
A.1 IMPUTATION
Imputation, the task of filling missing cormupted regions in data with plausible content, has been revolutionized by diffusion models across various modalities (Corneanu et al. 2024 Lemercier
et al., [025 Duan al., (024). Current imputation approaches follow three primary paradigms.
Palette (Shariia et al.) 2022) established conditioning on partially observed data by incorporating
known regions at each denoising timestep. RePaint (Lugmayr et al. 2022)) introduced a training-
method that leverages pretrained unconditional models by resampling known regions while generat-
ing content only for masked areas. Bridge- methods (Liu et al. 2023 Yue al.) 2023) Albergo
it al. 2023) design specialized diffusion processes between original and masked data distributions, requiring models trained to condition directly on masked inputs. DiffusionPDEE (Huang et al. 2024)
introduces diffusion model to solve PDEs under partial observation by learning the joint distribu- tion of coefficient and solution spaces. A critical limitation shared by all these approaches is their reliance on complete, unmasked data during training to learn the underlying data distribution before performing inference-time imputation on partially observed inputs. This assumption fundamentally conflicts with scientific applications where training data itself consists only of partial observations.
A.2 GENERATIVE MODELING WITH MISSING DATA
Deep generative models tackle missing data through various approaches, including VAE-based
methods (Ipsen et al. 2020 Ma et al. 2020) and GAN-based methods (Li et al.) 2019 Li & Marlin
2020). Some diffusion-based (Ouyang et al. 2023 Daras et al.] 2023 Dai et al.] 2024] Simkus &
Gutmann 2025) generative models generate clean samples from missing data, though they rely on
heuristic intuition and lack rigorous convergence analysis. DiffLight (Chen et al. (024a) leverages a partial rewards conditioned diffusion model to prevent missing rewards from interfering with the learning process. Zhang et al. (2025a) presents a theoretically sound framework combining diffusion models with EM algorithm for imputation, but its requirement of multiple complete model retraining cycles limits its scalability to complex and large datasets. More recently. Givens et al. (2025) pro- posed score matching with missing data, providing theoretical guarantees but facing computational scalability challenges due to high complexity in their importance weighting method and requiring auxiliary network training for their variational approach. Their experimental validation is limited to low-dimensional synthetic data and simple graphical models, raising questions about scalability to high-dimensional wooldworld scenarios.
B DATA MATCHING DIFFUSION MODELS
Diffusion models (Sogg& & Ermon 2019 Song et al. 2020 Ho et al., 2020) generate samples from a target data distribution by defining a forward process that gradually adds noise to data ap ~ Po according to the stochastic differential equation:
dx;- (t) + g())dw,, To ~ Po=
(15)
where w. E I2d represents standard Brownian motion, f(t) d g²
daf 2 dog as of
dt and at, are predefined time-dependent functions. This process has the analytical solution Pt (xx 20) = N (æx agTa, of? I).
dt
de
The reverse process generates samples by integrating backward from noise to data using: dæ,
f(t)a, - 1 - (1)V (o:)) ET Pr (2T))
(16)
dt
Since the terminal distribution (xT) becomes approximately Gaussian through appropriate pa- rameter choices, sampling from and reversing the process yields samples from Po- The score function V, = log (2;;) is computationally intractable but can be estimated using neural networks via noise matching and data matching approaches (Zheng et al. 2023): :
T noise (0) Er,
(ttt) ||ea (æt, t) <l(²]
E (x t)
o.
V logpt (x))
(17a)
1
+ of v log (a:), (17b)
T data (0) = =Er wott) ||00 (x), t) - x
s (zt, t) =
an
at
14

---

Published as a conference paper at ICLR 2026
where 2g Cg + d,e with . (.,
D. and w(t) is a loss weight function. By Tweedie's
formula (Efron 2011 Kim & Ye, 2021 Chung et al. 2022), we also have zq (æ,, t) = E(zo xt|-
C METHOD COMPARISON AND THEORETICAL ANALYSIS
C.I DIFFUSION PROBABILISTIC FIELDS
Several prior studies have proposed training generative models using field representations (Du al, Dupont et al. 2022 Zhuang et al. 2023). Similarly, our approach trains model to predict 2Dg Tobs.t, MD and applies the loss function to several randomly selected points across the entire set of coordinates. However, three main differences distinguish
202
approach from DPF: (1) Loss objective: Our method uses data matching while DPF uses noise matching, providing more flexible input requirements; (2) Model architecture: DPF requires both context and query inputs, whereas our method trains using only context, leading to different optimal solutions; (3) Theoretical founda- tion: DPF relies heuristic designs without convergence guarantees, while our approach provides rigorous theoretical analysis. As shown in Appendix [D] DPF's optimal solution depends on specific context-query mask combinations and may predict values differing from target predictions, whereas our method guarantees convergence to desired objectives through sufficient context mask sampling.
C.2 AMBIENT DIFFUSION
While ambient diffusion (Daras et al. 2023) shares the fundamental principle of our ap- proach-incorporating masks during training to predict clean data several critical distinctions emerge upon closer examination. First, ambient diffusion lacks thcoretical analysis of how differ- ent mask distributions affect the learning dynamics, whereas our work provides rigorous charac- terization through Theorem parts (ii) and (iii.. The most significant difference lies in the sampling methodology. Ambient diffusion employs fixed mask sampling strategy and directly approximates E |en Max @ Bobs,tt Mex] to approximate E Jro aaas,,,, M). As we demonstrate in Theorem 2 this approximation introduces distribution gap proportional to the variance of information provided by the conditioning terms, leading to suboptimal sample quality. In contrast, our method leverages the Martingale convergence theorem approximate the true conditional expectation, providing the- oretical convergence guarantees and eliminating the distribution gap inherent in the ambient diffu- sion approach.
C.3 SCORE MATCHING WITH MISSING DATA
Our approach differs from this prior work (Givens et al. 2025 in key way: we employ data matching rather than score matching due to its superior input flexibility and computational efficiency. In data matching, the optimal solution is E [xo Lobs,t, M). which requires only the observed data Tobs, and mask M as inputs. By contrast, noise matching and velocity matching require computing the score term Vaust. log P: (Irumobs,t è Dabs, M), which depends the unobserved data F that is unavailable during training. To address this limitation, the score matching approach relies on Monte Carlo approximation, significantly increasing computational cost. Additionally, prior work (Zhou et al. 2024) has shown that data matching achieves superior performance compared to noise matching for PDE solution generation. Therefore, focus exclusively on the data matching framework.
C.4 EXTENSION TO OTHER SCORE-BASED MODELS
Our approach can be extended to other score- acore-based generative models through established theoret- ical connections. The equivalence between diffusion models and flow matching (Albergo et al.) 2023 as well as between diffusion models and Bayesian flow networks (Xue et al.] 2024), pro- vides foundation for this extension. Through Tweedie's formula (Efron) 2011). we can es- tablish the connection between the score function Va log Pt (æx) and the conditional expectation E [zu 2..]. Leveraging this connection via the parameterization trick, we can train model to learn [on Lphs,t .M] using our proposed method, which can then be converted to output the score function for the denoising process (Zhou et al. 2024
15

---

Published as a conference paper at ICLR 2026
C.5 COMPARISON TO MASKED SELF-SUPERVISED LEARNING
Our context-query partitioning strategy is conceptually related masked signal modeling, such as Masked Autoencoders (MAEs) (He et al. 2022). However, our approach differs fundamentally in its problem setting, data assumptions, and theoretical goals.
Data Completeness Assumption (Most Critical): MAEs and related self-supervised methods train complete data that is artificially masked. The full ground truth is always available during training. Our framework is designed for a different, and common, scientific scenario: the training data is inherently incomplete (e.g., from sparse sensors cloud occlusion). No complete ground truth samples exist in our training dataset.
Primary Objective: MAE uses masking as a pretext task to learn robust representations for downstream applications (Feichtenhofer et al. 2022). Our goal is to learn the complete gener- ative distribution Pdaa(Zo) from these partial observations to perform accurate imputation of the true physical fields.
Masking Strategy and Theory: MAE's random masking is an empirical choice for representa- tion learning. Our strategic context-query partitioning a direct consequence of our theoretical analysis in Theorem. []] It is specifically designed to solve the core challenge of our setting: how to ensure that dimensions permanently missing in the training set still receive positive query probability (P((aay);
Max) 0) and meaningful gradient updates. As shown in Theorem. | without this, the model learns arbitrary values for these unobserved regions. This is problem MAE does not encounter, it always has access the complete ground truth for its loss calculation.
C.6 METHOО COMPARISON AND THEORETICAL ANALYSIS
We provide additional method comparisons Tab. 4 and summarize the key feature comparisons with the related methods, contrasting our approach with three existing methods: (1) Ambi-
entDiff (Daras et al. 2023), (2) Diffusion Probabilistic Fields (DPF) (Zhuang et al. 2023). and (3)
Score Matching with Missing Data (Givens al. 2025
Table 4: Comparison of the most related methods.
Aspect
Ours
Amblent Diffusion
DPF
MissDiff
Truining Objective
Data matching
Data matching
Noise matching
Noise matching
Model Input
Context only
sampling
context query
Masked tabular
Mask Usage
Hidden during training
Hidden during training
Provided to model
Provided model
Mask Sampling
Random subsets: M.. M
Fixed distribution
Both Mc. and May Incurrect
M.. Ma. M
Expectation Approx.
Ensemble: a E, mp(t,ctx) ctx)
Direct: E[py | ctx]
Direct: E/zg obs]
Theoretical Guarantees
/ Convergence proofs (Thm.
2
x Lacks rigorons analysis
Heuristic design
x Lacks rigorous ana lllssi
Distribution Gap
Minimized via ensemble
Gap x Var(conditionig]
Not addressed
x Not addressed
Learning Dynamics
Gradicat scaling analysis
No analysis
No anslysis
analysis
D METHOD COMPARISON: DIFFUSION PROBABILISTIC FIELDS
We summarize the training and sampling algorithms for diffusion probabilistic fields (DPF) Zhuang et al. 2023) in Alg. 3 We cannot directly adopt DPF for dynamic completion tasks because we lack access to the ground truth value of qry during training. However, we can still compare several high-level ideas with those in DPF.
16

---

Published as a conference paper at ICLR 2026
Algorithm DPF training process (Zuuagg
Algorithm 3 DPF sampling process (Zhuang et al. 2023
et al. 2023
Sample Max S May
repeat
2: av N(O,I)
4ig ~ Pdeas Unilorm (0. 1)
3: cix MMax. Max O OFrl
Eexx (0, Egry ~ NO,
4: qry = Mayy Mary @ @FT
4:
Sample Max: Mary
5: for T..... T...., do
5:
ctX Maxx Max @ ta + ceccl]
6: -
(0,
6:
qry Mary May @ (a,a0o + dxEegy)
7: ctx M Max @ #t-1
7:
Optimize the loss function
8: qry [May: May ④ 0--1
Mayy @ (ea(f, crx, Enpy) ||
9: end for
until converged
10: return sample values evaluated on May
There are three main difference between our proposed method and DPF:
1. Diffusion loss objective: our methods use data matching while DPF uses noise matching. Using data matching predict conditional expectation has more flexible requirement model input (see Sec. [C6 for details).
2. The DPF method suggests we should take Mary to be the full observed mask, which impos- sible to implement in our task.
3. DPF trains a model that takes both and qry as inputs. In contrast, our method trains a model using only ctx input. This difference leads the model to converge to different optimal solution.
In the following analysis, we will assume that we have the fully observed sample during the training session. We will analyze the output of the model optimized by Alg. and demonstrate that it does not yield the desired solution for the denoising process.
We reparameterize the loss function
(0) = Mgy @ (æe(t, ctx, qry) - xo) 1|2
(18)
When optimized,
Mgr O (agag (t, ctx, qry) + o: (t, ctx, qry) - qry(1))
(19)
In the following, we will analyze the optimal solution given by equation 18 The conditional expec- tation E [xo ctx, qry] minimizes the expected squared error (Bishop Nasrabadi 2006). Hence, we have
Mar @ (t, ctx, qry) Mnr [20 ctx, qry]
(20)
For simplicity, we consider simple distribution æp ~ Pdata N(u,, E). The noisy observations are defined as: Zatx
04 + 0feex: Eupy
0r TxEapy. We have joint Gaussian distribution:
20
Σ
E
Zax
~ ~
04
a
+ of
a
(21)
Euryy
a
a
of
of
Let y denote the observed entries selected by the masks:
=
Zetx
where is the selection
Eqry
matrix that extracts the masked entries from Zax and Zyry- By the property of linear transformation of multivariate Gaussian, suppose we have a joint Gaussian distribution the form:
x
A
Σ Drr
Σ,z
(22)
[2
Σ =z.
Σ
then the joint distribution of æo and y is given by
xg Sz
N
ur
Σ
ETSST
(23)
Sp=
SE
ST
Using the above property, the joint distribution of xo and is then:
36g
36p
S
Zetx
(24a)
Zary
W
8 (1a.s1
μ (12 (1,((12 1
a.
Σxn a1, Σ)
s
a4 (1, & Σ) E)S ST (a
& Σ) of1) ((1)s? ST
(24b)
17

---

Published as a conference paper at ICLR 2026
Thus, the conditional expectation is:
Do y] μ or (12 E) ST [S (a (121) & Σ) & + (y S p) (25)
Therefore, DPF's optimal solution depends on the specific context and query masks chosen, and may predict values that differ from target predictions. In contrast, our proposed method has theoretical guarantee: with sufficient sampling of context masks, the model's output will converge to the desired objective (see Theorem
E MULTI-STEP SAMPLING
In the following, will consider two specific settings of the mask M and diffusion time t, and use these constructions to approximate E [mo Øt; Lobs, M).
Diffusion expectation approximation. We randomly generate multiple masks {M md Mmd follows the same marginal distribution as Max in the training process, but not necessarily being a subset of M and take the average across all samples yields:
(6) k where
E (20 at
Σ K
xo
.M
(N) rad
De, M
(K) ml
(26)
This Monte Carlo estimation demonstrates that our model, despite being trained exclusively on masked data, can recover the full data distribution. The averaging process allows to obtain the same distributional modeling capability standard diffusion models trained on complete datasets.
Imputation expectation approximation. Given partially observed samples Lobs we apply the forward diffusion process to small timestamp t = 8 and generate random masks {M (K) K approximate:
S M
ctx Ik=t
E (20 Wobss M]
~ E (aro Bohs,, 6, M]
(k) @ Tots, M(k) (k) Σ wa (6, MMM
(27)
K
The optimal denoiser E [wo) Tr, Tobs: M] requires the expectation of xo conditional on all three information sources: the noisy state æt, the clean observations Lbbs, and the observation mask M. Our approach decomposes this complex conditioning into two manageable components: the diffir- sion expectation [an æx captures the denoising information from the current noisy state, while the imputation expectation E (rr aohss M] incorporates the structural information from the ob- served values and their locations. We then heuristically combine these two sources of information through a weighted average:
zig (1,, De, Robs, M) = I. [mo æie Fobs, M] R w, E. (mo axe] + (1 - -wy) E [20 Dobs, M] w.EE,,. (2p (1, Mna æ4, Mraa]] + (1 w.) E EM.CM (wo (6, Max Bohs, Mcx)]
(28)
where wi is a monotonically increasing weight function that transitions from 1 the dif- fusion process progresses and is sufficiently small positive number.
Algorithm 4 Diffusion-based sampling for data imputation
Require:
partially observed data Lohc. mask M, trained model Fg
We present our proposed sampling algorithm in Alg. At the implementation level, we precom- pute the imputation expectation once before the denoising process begins. During each denoising step, we approximate the diffusion expectation by sampling a single random mask Mnd and mak- ing one model evaluation. This single-sample ap- proximation is analogous to using a batch size I in stochastic gradient descent. where accept the variance from using only one sample in ex- change for computational efficiency. Following
Ensure: imputed complete data Zo
1: initialize: TT (0,
2: for diffusion steps from do
Cunobs f (in .M. (Eq. 28
4: Cobs x- 3e, 20,
o,
5: Cfall t- M Eobs Eobs (1 M) O Eunobs
6: 2. +- DiffusionODE((, t, Vs. Efull)
7: end for
8: Zp <- M Ø Datn + (1 M) ④
return 2o
18

---

Published as a conference paper at ICLR 2026
Lugmayr al. (2022), our sampling procedure applies different denoising strategies observed and unobserved regions. For unobserved elements, we estimate the noise using our trained model, directly compute the noise from the known clean observations. This
while for observed elements, approach ensures that the observed values remain consistent with their true underlying data through- out the denoising process. A more advanced setting proposed in |Huang al. (2024) uses guided diffusion sampling that starts from Gaussian noise and iteratively denoises it while being guided by two loss terms: observation loss (matching sparse measurements) and PDE loss (satisfying the governing equation), ultimately generating complete solutions that are consistent with both the par- tial observations and the underlying physics. However, we did not implement this approach in our work, and combining our mask-based denoising strategy with physics-informed guided diffusion remains an interesting direction for future research.
F PROOFS
This section provides detailed mathematical proofs for the main theoretical results presented in the paper. We begin by establishing key assumptions and notation that will be used throughout the proofs.
Assumption 1 (Uncorrelated decomposition). Decompose the model output as De (f, ctx) = [a +b (ctx)+ Cbias (ctx). Given context ctx, the following three components are mutually uncorrelated:
data component: Elro ctx] - [F[E(aro ctx]]]
bias component: b(ctx) [[((tx)]].
random error: e(ctx).
We also assume that ((ctx()) and e(ctx(i)) are independent for For notation simplicity, we denote
obs M @ Bobsst M]
(29a)
ctx Max O Febs,t Mex.]
(29b)
qr Mary ① Fobs,t M
(29c)
when is clear from the context.
F.1 ANALYSIS OF MODEL OUTPUTS UNDER OPTIMAL LOSS CONDITIONS
We begin the proof with a foundational lemma that establishes the key relationship between the optimal model output and the conditional expectations.
Lemma 1 (Optimal Function for Element-wise Weighted MSE). Let X, Y, and Z be random vectors in Id such that the relevant second moments are finite. Let g Space(Y) Rd be deterministic funciiioo and let the objective function L(g) be defined as
I(g) = F. [1] g Z x
(30)
where @ denotes the Hadamard (element-wise) product. If each component of the vector Z|Z c Z Y) is strictly positive almost surely, then the unique function g" that minimizes L(g) is given by
E
Z
c
X
Y
g'
(31)
E
Z
Z
Y
where the division is performed element-wise.
19

---

Published as a conference paper at ICLR 2026
Proof. The objective function L(g) can be decomposed by writing the squared Euclidean norm as a sum over its components.
T(g) E Z ④ (g(Y) - x]1
(32a)
(i(gi xx))²
E
(32b)
Σ (gi x.)
(32c)
The final step follows from the linearity of expectation. The total loss is sum of non-negative terms, so L(g) minimized if and only if each term in the summation is minimized independently.
Let (g,) RJ Z (Y.(Y) -X be the i-th term. The optimization problem is thus reduced to
finding the function g, that minimizes Lg for each component E [l.....d]. By the law total expectation, (g.)) can be written as:
L(gg)) = Ev (g) - x.) Y|]
(33)
The outer expectation is minimized by minimizing the inner conditional expectation for any given realization y from the space of Y. For fixed y. let v, = gzy be a deterministic scalar. The inner expectation becomes:
E
[ (v. x,) Y y) = [Z 2v, X;
x?)
Y v]
(34)
Applying the linearity of conditional expectation yields a quadratic function u,;: B Y y] 2v, 2,, Y = yو + E
Y
= y
(35)
This is a convex quadratic v;, since its leading coefficient E Z Y ya is strictly positive by hypothesis. The unique minimum is found by setting the derivative with respect to to zero:
20; E ? Y y)
2 X; Y y)
(36)
Solving for n; gives the optimal value for the component function at y:
E x Y = yl
(37)
y]
This establishes the optimal form for each component g (y) of the function gy)).. Since this holds for all y, the optimal function 9 for the i-th component is:
E ? X;
g.
(38)
E Y]
Assembling the components for I,... d into single vector equation gives the expression for the optimal function g":
E Z Z Z
g
(39)
Z ④ Z Y
where the division is understood to be dlmeent-wise.
Having established the fundamental relationship between optimal model outputs and query proba- bilities in the lemma, now proceed prove main theorem.
Theorem 1 (Optimal solution under context masking without query information). Let z be the optimal solution by minimizing the loss equation Under the conditional independence of masks and data, we have the following results:
(i) Optimal solution:
The optimal solution given by
E ( M. ① Dotn.,t Man.. (we (t, Men. 3 Robr,ts Mau)); an arbitrary value,
P((May); = M. > 0 P((M = M =0
(2)
where indicates the i-th entry of the vector. Specially, given the context mask Max if the union of all possible query mask M supports covers all spatial dimensions, we have
Do (1,, Ma Bobs,t, Man) E [20g
Max O Bobstt Max)
(3)
20

---

Published as a conference paper at ICLR 2026
(ii) Gradient magnitude scaling: The expected squared gradient magnitude with respect the ner- work output for dimension scales linearty with the query probability ps == P((Mar)) = Mex):
(
Oc (x0)):
)
E
4p,E ((xz):
(Zobs)
(Mary =
(4)
(iii) Parameter update frequency: The frequency of non-zero parameter updates for dimension is exactly ps:
Pidimension i contributes parameter update) = P((M =1 Mar) = Pi
(5)
Proof. Given the training algorithm in Alg. [ we have
0 arg min E au.. (M.. Map CM) ||| Mary @ (ro (1, ctx) - (obss) oobb))
(40a)
arg min E (Mar Ma CM)
(M
[| @ (aro (1, ctx) M ④ acn) [|
(40b)
arg min EE (Ma. CM) [| C xe (t, ctx) Mary + xo [
(40c)
(i) Optimal solution:
When optimized, given MMex S M, for any index i such that P((Mgry)
1 Max)= > 0 under the sampling distribution, we have
May ② To (t, ctx)
E Ma @ æu ctx, Mayy
Mar
E [xu ctx]
(41)
Thus, Mar O (ze ctx) E [xr ctx)) 0 for any Mar in the support of the sampling distribu- tion given MÆx.
Case When P((May ( 1 Mexx 0. By applying Lemma with X Cy. and g(-) = (t, -), the optimal solution is:
MMayy Y ctt,
(1, May @ May O æo ctx]
(42)
BB. Bay O Mary ctx)
Since MAry is a binary mask where (Mary): E {0, 1}. we have (May) (May)i ) = (May), Thus:
Mry æo ctx]
2z (tx
(43)
EE M|yy ctx]
For component where P((Mar)) ) = Mcx)) > O:
(= ctx]
(t,
(44a)
(gyy) )s ctx]
E[( [M ) ctx]
(44b)
P((Mar): = 1 Max)
E (M ctx]
(44c)
P((Mar): = Max)
= [[((w)); ctx, (Magy : = 1] E[((x)); ctx]
(44d)
(44e)
where the last equality holds because wo independent of Mar given ctx.
Case 2: When P((M = 1 Max) 0. In this case, (Myyy) = 0 almost surely given Max- The contribution of index to the loss function is:
E|[ (zo etx)); - (May)a ctx]
= |0 (zo(t, ctx)); 0- (woo) ctx]
(45a)
0
(45b)
Therefore, (1, ctx)) does not affect the loss function and can take any arbitrary value.
Thus, we have
[E [(xr); ctx),
P((Mpr); Max) 0
(xo (t, ctx)),
(46)
an arbitrary value,
P((May); = Max) = 0
21

---

Published as a conference paper at ICLR 2026
If the union of all possible query mask May supports covers all spatial dimensions, have
UM,, possible supp(May) = (1...., ( Thus, Vi, P((Mar)) = =) > and
zg (t, ctx) = E z ctx]
(47)
(ii) Gradient magnitude scaling: The gradient of the loss with respect to the i-th output component is:
dc M( ((xe); (2oots)))
(48)
((r));
Taking expectation over the sampling distribution:
( ac
((zz)))) )
)
E
E ( ); ((xro); (zobs)))
(49a)
4E [M gry) ((xr); (z)
(49b)
E ((x)): (z) (Mary)s= =
(49c)
Let C E ((xa): (Iota);) (Mary) ) 1]. Then:
(
Oc
| 11-4pe -
E
4p,0.
(50)
0(x));
establishing linear scaling with p;.
(iii) Parameter update frequency: At each training step, the gradient contribution from dimension
i is:
aC dC 0(xo);
(51)
00 ;
O(xr); 00
This contribution is non-zero if and only if (May); 1, which occurs with probability ps. There- fore:
P(dimension i contributes to parameter update) P((M = Maa))
(52)
F.2 PARTIALLY OBSERVED MASK CONVERGENCE THEOREM
Theorem 2 (Ensemble approximation convergence). Let E| obs] = E[vo Tobs M] be the ground truth conditional expectation and E[xo ctx] context-condittee conditioned expectation. The expected squared error between these quantities is:
E E(zn M O Dobs,t M be the
E ||xx
ctx]
E[zu
[bs]]] E[Var[æu ctx]] E[Var|æo obs]]
(11)
Consider a practical model with output wo (1, ctx) = Rf ctx] +b(txx) + Exias (ctx), where b(ctx) represents context-dependent deterministic bias and (ctx)(et)) is random with El(esirr] 0. Given the ensemble prediction in equation as:
Ax
1
Σ ro (t)(
(12)
K
the expected squared error between the ensemble prediction and ground truth is:
E
||ūx
E|zo
obss]]
EE
ctx]]
Bfaco
obs) + E(b(ctx)) ||
information gop
model bins
(13)
Var
[[ag data variance
R
ctx]
Var[b(ctx]
+ Var(Ebial
moutet variance
As K +
the ensemble converges to:
lim E [lAk
E[zo
obs]][?]
|
EE
ctx]
E[z0
obs] + E[b(ctx)]]1
(14)
K+000
22

---

Published as a conference paper at ICLR 2026
Proof. Define random variables μ [zu obs] and pe E [l|e- |ell"]-
E [anu ctx].
Our goal is to compute
Applying the Law of Total Variance, we have
Var (wo ctxj = E [Var [ato obs]
ctx]
+ Var [E [ obs]
ctx]
(53a)
E [Var [z0 obs]
ctx]
Var ja ctx]
(53b)
Rearranging:
Var [ee ctx]
Var (wo
ctx]
IE [Var (20
obs]
ctx]
(54)
Then:
E [ll Mella]
E [ [
Mel
2
ctx]
Law of total expectation
(55a)
E E [||
E
ctx]
|1
ctx
By definition of pe
(55b)
EE EE
|4
E
[E
[20
obs]
ctx]
|12 ctx
Tower rule
(55c)
E E
E
[x
ctx)
||
ctx]
By definition of μ (55d)
E Var
ctx]
Definition of conditional variance
(55e)
E [Var
ctx]
E [Var
(zg [bs] ctx]
equation 54
(55f)
E Var
[mo
ctx]
E [aar (zq obs]
Linearity of expectation
(55g)
Further define μa = vo (t, ctx E x
ctx( + (ctx + Chias (ctx(*)), where is the
systematic bias and ébins is the random error with IE' [Ebiasl 0. The ensemble average is:
K
K
K
K
Σ
Eizo
ctx ctx
+
1
Σ
b(ctx (1)) +
Σ c(ctx (1)
(56)
K
K
K
K
Computing the expectation:
K
E
E
14
(57a)
K
k=1
1
B[B[200
ctx + ((ctx(()) ((ctx(()))
(57b)
K
E
[E azo
ctx]] E[b(ctx)] + B[e(ctx)]
(57c)
= E3
[ (
ctx]] [[((ctx)]
(57d)
The bias of the ensemble estimator is:
Bias( = x[ - E[æo obs] E[E[ ctx]] - E[wo obs] + [[b(ctx)]
(58)
For the variance, decompose each component:
p4 [[ = E [00 ctx (+)] + ((ctx(()) + E(ctx - EE ctx]] - E[b(ctx)]
(59a)
E ED
ctx (x)] E[E[a0 ctx]]) + (b(ctx (+)) Elb(ctx))) + e(ctx (k)
(59b)
data
The variance of the ensemble average:
Var(ūx)
E
|| 11
(60a)
K
E
(60b)
K
K
||
E[ma]112
+2 >
(m:: Blpas..15 Elbez]>
(60c)
K2
23

---

Published as a conference paper at ICLR 2026
Since ctx(1), ..... ctx conditionally independent given Zobs, and M, the data variations and bias variations independent across different k. Additionally, under the assumption that e(ctx( and ((ctx()) are independent for *
Elp Elu.) py B(sI) = vi j
(61)
Therefore:
K R 1e [B]
Var(px)
(62a)
K
[B(aa ctx] [[
ctx]]) (b(ctx) [B[b(ctx)]) <(ctx))l²
(62b)
(Var|Elaa
ctx]] + Var[b(ctx)] + Var(e])
(62c)
K
Here, we used the independence between the three components to separate the variances. Combining bias and variance using the bias-variance decomposition:
E |
El obs]][]]
Bias I1 Var(px)
(63a)
= 1x0 ctx]] - E| obs] + E[b(ctx)]]]
1
(Var[E[zoo ctx]] Var [a(x))] Var(e])
(63b)
K
Taking the limit as K + oo:
lim +00
|lıx
E Doo
obs]] = -|[|E
[tx]]
E(zop
obs] + E[b(ctx)]]]
(64)
This establishes that:
The average bias E[b(ctx)) across all contexts is not reduced by ensemble averaging.
The variance of the context-dependent bias Var[b(ctx)] is reduced by factor of 1 1/K.
Both data variance and random error variance also reduced by 1/K.
The asymptotic error includes both the data bias and the model's average bias.
G DISCUSSION ON INAPPLICABILITY OF BASELINES
G.1 INAPPLICABILITY OF KNEWIMP
We discuss the inapplicability of the primary method proposed in Chen et al. (2024b), Kernelized Negative Entropy-regularize Wasserstein gradient flow Imputation (KnewImp), as a baseline for high-dimensional PDE dynamics task.
Principle. KnewImp is an approach explicitly designed for the imputation of numerical tabular datasets. The method reformulates the imputation problem within the framework of Wasserstein Gradient Flow (WGF). Its core contribution is to derive a closed- form. implementable imputation procedure by optimizing a novel Negative Entropy-Regulant producing Kernel Hilbert Space (RKHS).
(NER) cost functional within a Re-
The final imputation procedure involves simulating an ODE velocity field (x (joinl), .T) is defined using a kernel function:
dxx (miss)
2xX (joint, 7), where the
dr
u (joint). -) = E, (rmm .T) {+|V & (mia) log Pl X (oint)) )) KC (joint) X (Goin)) }
AV & (miss) KC (Join)
(joint) (Joint)
(65)
24

---

Published as a conference paper at ICLR 2026
This velocity field depends two components: (1) kernel function K., which the authors specify is a Radial Basis Function (RBF) kernel, ( X) exp(- ||x-X] and (2) an estimated score 2h
function V x ( log (x (juin)), which is trained separately using DSM (Song et al. 2020
Reason for Unsuitability. This method is mathematically and computationally infeasible for our task for two primary reasons:
Mathematical Infeasibility (Curse of Dimensionality): The entire derivation of the imple- mentable velocity field hinges on the use of an RBF kernel. This kernel's s computation is based on the squared Euclidean distance | X| between data points. Our PDE dynamics data has dimensionality of D 100 64 x 64 409, 600. In such an extremely high- high-dimensional space, the concept of Euclidean distance becomes meaningless; all data points tend to become equidistant from another. This "curse of dimensionality" would cause the RBF kernel to lose all discriminative power, rendering the velocity field calculation mathematically unstable and uninformative. The KnewImp method is fundamentally structured for low-dimensional tabular data, where distance metrics remain meaningful.
Prohibitive Computational Cost (VRAM and Time): The method's two-stage process scales intractably with dimension D.
The "Estimate" phase requires training a score network V (jmn) log ((X (jvin)) via DSM. Training any neural architecture (U-Net, Transformer, etc.) to model the score of D 409, 600 dimensional vector would require prohibitive amounts of VRAM just to store the activations and gradients for single batch.
The "Impute" phase requires simulating an ODE. Each step of this simulation necessi- tates a full computation of the velocity field. This computation involves a Monte Carlo estimation over the dataset, with each sample calculation requiring an expensive kernel evaluation in D-dimensional space.
This combination of a memory-intensive score network and a computationally--neesise kernel-based ODE simulation makes the Knewlmp approach computationally infeasible for our high-dimensional spatiotemporal task.
G.2 INAPPLICABILITY OF SCORE MATCHING WITH MISSING DATA
We discuss the inapplicability of the two primary methods proposed in Givens et al. 2025) as base- lines for our high-dimensional PDE generation task.
G.2.1 METHOD MARGINAL IMPORTANCE WEIGHTING (MARG-IW)
Principle. This approach (Algorithm 1) adapts the score matching objective to work with missing data by defining a marginal Fisher divergence. It then estimates the intractable marginal scores by applying importance weighting (IW) to approximate the high-dimensional integral over the missing coordinates using Monte Carlo samples.
Reason for Unsuitability. This method is unsuitable as it is not designed for high-dimensional data. The authors of the original paper explicitly state that the IW approach "will struggle in higher dimensional scenarios" and primarily demonstrate its efficacy in "lower dimensional set- tings". Given the dimensionality of our task, the variance and bias from the IW estimator would render the optimization intractable.
G.2.2 METHOD MARGINAL VARIATIONAL SCORE MATCHING (MARG-VAR)
Principle. This more complex approach (Algorithm 2) avoids direct IW estimation by first taking the gradient of the loss objective (a "gradient-first" approach). then introduces secondary vari- ational neural network (v')) to approximate the conditional expectations and covariances over the missing data. The training involves a nested optimization, where the main score model (sg) and the variational "helper" model (p')) updated iteratively.
25

---

Published as a conference paper at ICLR 2026
Reason for Unsuitability. This method is computationally infeasible for task due pro- hibitive VRAM and time costs. Our backbone model (analogous to se) already consumes ~00G of VRAM for a single barkwrdd pass with batch size of 8 on an 80GB A800 GPU. The Marg- Var algorithm would introduce, at minimum:
Dual Model Memory Cost:
The algorithm requires maintaining two large neural net- works-the primary score model 8o and the variational helper p/s, simultaneously in VRAM. This alone would exceed the 80GB capacity.
Peak Gradient Memory Cost: The gradient calculation for so (Eq. 10/11) is exceptionally complex. It requires computing expectations and covariances that involve components from both models, necessitating that the computational graphs of both networks are active for the joint gradient estimation. This leads to a peak VRAM usage far exceeding the simple sum of the two models.
Multiplied Training Time Cost: The algorithm employs nested optimization loop. For each gradient step of the main model 8g, the helper model p's must be trained for L steps (e.g., L 10 in the paper's experiments). This L-fold multiplication of an already lengthy training step makes the method impractical for large-scale generative modeling.
G.2.3 DIMENSIONALITY MISMATCH
The core issue is the extreme discrepancy in data dimensionality. The experiments in Givens et al. (025) are conducted tasks with dimensions of 10 (Gaussian), up to 50 Non-Gaassian), 100 (S&P 100), and 106 (Yeast). Our PDE dataset has a dimensionality of 100 x 64 x 64 = 409, 600. This is more than three orders of magnitude larger than the highest-dimensiooal task (106-dim) on which the Marg- Var method was validated. The computational and statistical challenges of this scale are not addressed by these methods.
H SUPPLEMENTARY EXPERIMENTS
H.1 DATASET SETTINGS
Shallow Water and Advection datasets. We evaluate our approach on two fundamental geophyy ical PDE systems with distinct characteristies:
Shallow Water: d.u fo 90gh, 0, fu -u 0,h HH + 0,v)
(66a)
Advection: Oru(t, x) + 38, x) =)
(66b)
where for the shallow water system, u and v represent velocity components, denotes height field, is the Coriolis parameter, is gravitational acceleration, and H is mean depth. For the advection equation, u represents a scalar field being transported and / is the advection velocity. We gener- ate synthetic solutions by randomly sampling physical parameters, along with randomized initial conditions to ensure dataset diversity. Each dataset contains 5k training, 1k validation, and 1K test samples with 32 x 32 spatial resolution and 50 temporal frames. We generate synthetic solutions by randomly sampling physical parameters and initial conditions. For evaluation, we employ comple- mentary metrics comprehensively assess reconstruction quality across different physical aspects. For the Shallow Water dataset, we evaluate PDE feasibility loss, which measures how well the re- constructed solutions satisfy the underlying shallow water equations by computing the residual error when the reconstructed fields are substituted into the governing PDEs. For the Advection dataset, we reconstruct complete initial conditions from partial observations using our trained model, then propagate these reconstructions forward using traditional PDE solvers (finite difference schemes) to generate temporal sequences. Performance is measured by computing MSE between our recon- structed solutions and ground truth sequences over all 50 time steps. This dual evaluation strategy demonstrates our method's superiority across multiple physically meaningful criteria.
Navier-Stokes dataset. for isotropic turbulence. The governing equations are:
We evaluate approach on the incompressible Navier-Stokes equations
d,u + (u V)u = -Pp Vp+uv2u
(67a)
V vuu
(67b)
26

---

Published as a conference paper at ICLR 2026
where (u, v) represents the velocity field, is pressure, is kinematic viscosity, and f is exter- nal forcing. The data are generated using either pseudo-spectral solvers with 4th-order Runge- Kutta higher-order Finite Volume IMEX methods. Initial conditions with varying peak wavenumbers eventually evolve to exhibit the Kolmogorov energy cascade. The dataset contains 1152 samples with spatial resolution of 64 64 and temporal sequences of 100 frames. For evaluation, the model generates complete field reconstructions. Performance is measured by computing MSE be- tween reconstructed fields and ground truth across all spatial locations and time steps.
ERAS dataset. We evaluate approach on the ERAS reanalysis dataset, which provides com- prehensive atmospheric and surface meteorological variables. ERA5 represents the fifth generation atmospheric reanalysis from the European Centre for Medium- combining model data with observations through data assimilation to produce a globally com- plete and consistent dataset. We utilize year of hourly data sampled at spatial resolution of 103 x 120 grid points in -niguudelooggiude coordinates, creating temporal sequences of dimen- sion (365 24, 103, 120). The data segmented into 3-hour windows, and only select 20% / 10% 1% observed pixels for each sample. Our experiments incorporate nine essential atmospheric variables. These variables capture both surface conditions and atmospheric dynamics critical for weather prediction tasks. Performance is evaluated by computing reconstruction errors across all spatial locations and temporal frames.
Weather Forecasts (ECMWF),
H.2 MASK SELECTION IMPLEMENTATION
We now present concrete algorithmic implementations of our strategic context-query partitioning framework for both pixel-level and block-wise missing patterns. Algorithm deetails the pixel-level procedure, where each observed location in the observation mask M is independently selected as context or query through Bernoulli sampling with ratios Tex and Tary, ensuring that Max M and Mgr M. For the more realistic block-wise scenario depicted in Algorithm [6] we operate on spatial blocks. This bock- sampling preserves spatial continuity while guaranteeing that every observable dimension maintains a positive probability of being selected as query, directly implementing the principle from Theorem I
Both procedures maintain the crucial property that P((May); Mcx) > for unobserved dimensions, thereby enabling effective learning from incomplete training data.
Algorithm 5 Pixel-Level Context-Query Partitioning
Require: Observation mask M E {0, 1)d, context ratio Tax E (0, 1), query ratio Tary E (0, Ensure: Context mask Max. query mask Mar
1: Initialize Max + O, Mgr +-
2: for each spatial index {1,.... d} do
3: if M, = then
4: Sample
Uniform(0,,
5: if u < Tax then
6:
(Max):
7:
end if
8:
Sample u Uniform(0, 1)
9:
if v < Tay then
10:
(M <-
11: end if
12: end if
end for
14:
return M Mgr
27

---

Published as a conference paper at ICLR 2026
Algorithm 6 Block-Wise Context-Query Partitioning (Integer-based)
Require: Observation mask grid Mgrid {0. 10 (e.g.. observed blocks out 9), integer hex (number of context blocks, e.g.. 4). integer Kupy (number of query blocks, e.g., 1)
Ensure: Context mask Max query mask Mar
Bobs < {bl (Mgnaa) =
Sample Bexx S Bobs uniformly with Bex = katx
Sample Bary Bobs uniformly with [aryl Kary
Max < BlocksTolke ToMask (Bex)
Mar f Blocks
BlocksToMask ToMask (Byry)
return Max: Mar
H.3 ANALYSIS OF MISSDIFF BASELINE AND DATA MATCHING ADAPTATION
A notable aspect of our experimental setup is the adaptation of the MissDiff baseline (Ouyang al.) 2023) from its original noise matching objective to data matching framework. This modification was empirically necessary, as the original objective proved ineffective in our experimental context. This adaptation facilitates a meaningful and fair comparison by ensuring the baseline can operate effectively our challenging datasets.
Initial failure of the noise matching objective. initial evaluation of MissDiff with its original noise matching objective showed that the training loss failed to decrease at all in experimental setting. This failure is attributed to a fundamental difference between the data domains: the tabular data used in the original MissDiff paper and the spatiotemporal data used in our work.
Tabular data (original MissDiff domain): The MissDiff paper focused on tabular data with relatively moderate missing ratios. In this context, each entry often represents an independent feature. Missing one entry means completely losing information about that specific feature.
Spatiotemporal data (our domain): Our PDE datasets involve spatiotemporal fields (e.g., images/videos) characterized by much higher missing data ratios (down 1% observed data). Critically, in these physical fields, a missing pixel does not represent the loss of an independent feature. Due to the inherent spatial smoothness and continuity of physical systems, neighboring pixels carry highly correlated information.
Implications for diffusion objectives. This fundamental data difference has profound implica- tions for the suitability of noise matching versus data matching:
Data matching (our adaptation): This objective (predicting 2o) can effectively leverage the spatial smoothness priors. Even from sparse observations, the model can learn to interpolate and predict reasonable values for missing regions by exploiting the correlated context.
Noise matching (original MissDifT): This objective (predicting e) requires the model to pre- dict fine-grained noise patterns. This task demands much denser observations to capture the necessary local structure. At extreme sparsity (e.g.. 1% observed), the noise prediction task becomes ill-posed. There is simply insufficient local context to distinguish signal from noise, making the learning target ambiguous.
Our empirical findings show that in our setting, the original noise matching objective led to MissDiff completely failing to learn (c.g., outputting all zeros). The adaptation a data matching framework allows MissDiff to produce meaningful predictions by leveraging the smoothness priors inherent in physical systems. Therefore, this modification was essential for valid and fair comparison. With- out this adaptation, MissDiff would be unable to generate meaningful results in our experimental scenarios, rendering the comparison ineffective.
H.4 ABLATION STUDY
We conduct ablation studies to validate the effectiveness of key components in our proposed method.
28

---

Published as a conference paper at ICLR 2026
Test-time gap introduced by replacing Max with M. Our sampling procedure requires mul- tiple context masks Max to estimate E |0 Lobs, M * xe M + Lt, M (k)
This ablation study compares our method against directly computing E [x0 Tobs,e, M) using ze (t, MM @ æt, M). The direct approach creates distributional mismatch: during training, the model's input mask follows the distribution of Max, but during sampling, the input becomes M. This mismatch degrades model performance. Tab. 5 presents experimental results comparing both methods.
Table 5: Performance comparison of two approaches: (1) imputation with multiple time sampling of Max followed by ensemble prediction (Theorem versus (2) directly using M as Max. which creates a distributional mismatch between training and testing inputs.
Shallow Water
Advection
Method
80%
30%
80%
30%
Ours (M)
2.3983 + 0.7880
2.6717 + 1.4731
0.1320 + 0.0155
0.1655 + 0.0537
Ours (Max)
0.1878 + 0.0054
0.7379 + 0.1101
0.1035 + 0.0008
0.1189 + 0.0069
Backbone architecture. To demonstrate the generalizability of our method across different archi- tectures, evaluate both our proposed approach and baseline methods using two distinct back- bones: Karras UNet (Karras et al.] 2024) and Fourier Neural Operator (FNO) (Li et al., (2020). For the FNO implementation, we concatenate diffusion time embeddings along the channel dimen- sion. Results are presented in Tab. O Our findings show that U-Net and FNO achieve comparable performance on the Shallow Water and Advection datasets, while U-Net outperforms FNO on the Navier-Stokes and ERA5 datasets, where FNO fails to generate reasonable samples.
Table 6: Performance comparison across backbone architectures. Results for our method and base- lines using Karras UNet (Karras et al. 2024) and FNO (Li et 2020 backbones across two PDE datasets.
Shallow Water
Advection
Method
Backbone
80%
30%
80%
30%
UNet
0.3963 + 0.0617
1.2570
+: 0.2146
0.1030 +: 0.0004
0.1
+ 0.0096
MissDiff
FNO
0.2917 = 0.1683
0.7525
+ 0.1529
0.1375
+ 0.0063
0.4816
+ 0.0187
Ours
UNet
0.3279 - 0.0655
0.9292
0.1963
0.1035
0.0008
0.1189
0.0069
FNO
0.1869 - 0.0015
0.7379 + 0.1101
0.1240 + 0.0040
0.3527
0.0620
Context and query mask ratio selection. We conduct an ablation study examining how different choices of context and query mask ratios affect model performance. The results presented in Table We evaluate ratios ranging from 0.5 to 1.0 to understand the trade-offs between information availability and parameter update frequency identified in our theoretical analysis. As expected from our theoretical framework, intermediate ratios (0.5-0.9) achieve optimal performance by balancing the information gap and parameter update frequency trade-offs. Notably, when both context and query ratios are set to 100%, our proposed method reduces to the MissDiff baseline, providing direct comparison point that validates our experimental setup.
Optimal denoiser approximation. We approximate the optimal denoiser E [Co Øe, Tobs, M] through a weighted combination of diffusion expectation [ æ and imputation expectation E [xa Lobs. MM] using empirical weight Wa (equation 28). We investigate different weighting strate- gies to understand their impact on reconstruction quality during the multi-step generation process (200 steps). The results can be seen in Tab.
29

---

Published as a conference paper at ICLR 2026
Table
Performance comparison of context and query mask ratio.
Navier-Stokes
Context Ratio
Query Ratio
80%
60%
20%
50%
50%
0.2383
0.5338
2.0924
70%
70%
0.2076
0.5336
2.0336
90%
90%
0.2252
0.5251
2.1309
70%
100%
0.2103
0.5276
2.1178
100%
100%
0.2444
0.7023
2.5599
Table 8: Impact of weighting strategies optimal denoiser approximation.
Navier-Stokes
Method
80%
60%
20%
w/o We
0.2334 + 0.1155
0.5649 + 0.0329
3.3820 + 0.1704
0.2331
0.0117
0.5633 + 0.0332
3.1881
+ 0.2170
W/ t2
0.2334 0.0116
0.5647 0.0333
3.3557
0.1973
Influence of ensemble size K. Tab. shows that increasing K consistently improves performance (see Theorem. 2). We use K 10 by default to balance efficiency and accuracy.
Training cost: Our training procedure has comparable computational cost to baseline diffusion methods Diff, AmbientDiff) since the network architecture, input/output dimensions, and num- ber of training steps the same. The main difference is in context-query partitioning strategy during training, which adds negligible overhead.
Inference cost: The additional computational cost comes from sampling:
For single sample generation (common in scientific applications): The K forward passes can executed in parallel since they are independent. Wall-clock time increases sub-linearly with A800 GPU, K 10 requires 3.36x the time of 1 for 50-frame 32 x 32 sequences, and 8.32x x for 100-frame 64 64 sequences. The overhead depends hardware parallelization efficiency and batch size.
K rather than K-fold. For example,
For batch generation of multiple samples: The computational cost scales approximately K times compared to baselines. This represents a fundamental trade-off: our method enables learning from realistically incomplete data, necessity in many scientific domains where com- plete measurements are physically impossible.
Table 9: Impact of ensemble size K Navier-Stoken imputation. Errors decrease with larger K but with diminishing returns. Time cost is measured for single-sample forward passes on a single GPU.
Navier-Stokes (x 10
Ensemble size
Time Ratio
80%
60%
20%
K = 1
0.2239
0.5652
2.1446
1.00x
K = 2
0,2147
0.5475
2.0822
1.81 x
K 3
0.2119
0.5418
2.0640
2.65x
= 5
0.2094
0.5371
2.0462
4.26x
K 10
0.2076
0.5337
2.0343
8.32x
K 20
0.2068
0.5320
2.0277
16,48x
K 50
0.2062
0.5308
2.0240
41.06x
30

---

Published as a conference paper at ICLR 2026
H.S COMPLETE RESULTS
We provide the complete experimental results, including standard deviations, to demonstrate the statistical significance and variance of our findings.
Table 10: Performance comparison on PDE imputation tasks. Each sample represents a temporal sequence of 50 frames, each with 32 x 32 spatial resolution. Results show the MSE between the reconstructed and the ground truth solutions from the PDE solver, averaged over all timesteps.
Shallow Water (feasibility loss, x x00=*)
Advection (simulation MSE, x10-1)
Method
80%
30%
80%
30%
Temporal Consistency
3.0248
4.2742
0,5097
0.6911
Fast Marching
2.5931
8.8631
0,2127
0.5222
Navier-Stokes
0.7045
2,8244
0,1350
0,4805
MissDiff
0.2917 + 0.1683
0.7527 + 0.1530
0.1030 + 0.004
0.1197 + 0.0096
AmbienDiff
0.1927 + 0.0050
0.7504 + 0.1119
0.1039 + 0,0009
0.1219 + 0.0075
step
0.1878 + 0.0054
0.7379 + 0.1101
0.1035 + 0.000K
0.1189 + 0.0069
Ours
200 steps
0.1869 + 0.0015
0.7502 + 0.1120
0.1037
+ 0.0009
0.1231
+ 0.0109
H.6 VISUALIZATION OF GENERATED SAMPLES
0
1-5
T- 10
t- 16
21
2-
t- 32
t 38
1-43
t- 49
Figure 4: Imputed results on 2D Shallow Water dataset where 30% of the original data points are available for training.
1-5
1-10 10
t-
16
21
2-
3-
t- 38
T-43
t- 1-49
Figure 5: Imputed results on 2D Advection dataset where 30% of the original data points are avail- able for training.
I LIMITATIONS AND FUTURE WORK
Our work represents a first step towards systematically incorporating mask distribution priors into the training of generative models for incomplete data. A primary assumption in our current frame- the data æn-
work is that the mask distribution Pracc (M) known a priori and is independent However, in certain real-world scenarios, the missingness mechanism can be data-dependent (e.g., a weather station failing due to the direct impact typhoon it is measuring) or follow complex patterns that unknown. Our current methodology does not explicitly address these more com- plex cases. We believe that extending this framework to handle unknown or data-dependent mask distributions is significant and important direction for future research.
31

---

Published as a conference paper at ICLR 2026
T-0
T-16
T-33
T=49
T-66
T=82
T-99
Original Data
Imputed Data
T=0
T=16
T=33
T=49
T=66
T=82
T=99
Original Data
Imputed Data
Figure 6: Sample imputation results on 2D Navier-Stokes dataset where 80% of the original data points are available for training.
T=33
T=66
T=99
Original Data
Imputed Data
T=0
T=16
T=33
T=49
T=66
T= 82
T= -9
Original
Data
Imputed Data
Figure 7:
Sample imputation results 2D Navier-Stokes dataset where 60% of the original data points are available for training.
On the theoretical front, our analysis provides guarantees for the asymptotic convergence of our training paradigm. We acknowledge that this analysis does not extend to a non-asymptotic regime. A more comprehensive theoretical investigation, such deriving finite sample complexity bounds or formally quantifying the approximation error introduced by the neural network architecture, is considerably challenging. Such analysis would need account for the complex interplay between
32

---

Published as a conference paper at ICLR 2026
T-0
T-16
T-33
T-49
T-66
T=82
T-99
Original Data
Imputed Data
T=0
T=16
T=33
T=49
T=66
T=82
T=99
Original
Data
Imputed Data
Figure 8: Sample imputation results on 2D Navier-Stoker dataset where 20% of the original data points are available for training.
T=0
T=16
T=33
T=49
T=66
T=82
T=99
Ground Truth
Ours
Miss
Figure 9: Comparison of original and imputed data from the Navier-Stokes dataset (one missing block). Each sample consists of 100 frames at 64 x 64 resolution.
the diffuuson process, the context-query sampling strategy, and the function approximator. We leave this rigorous theoretical extension an important open problem for future work.
LLM USAGE STATEMENT
We used large language models (Claude) to assist with manuscript preparation in the following ca- pacities: (1) improving the clarity and grammatical correctness of our writing through proofreading
33

---

Published as a conference paper at ICLR 2026
T=0
T=16
T=33
T=49
T=66
T=82
T=99
Ground Truth
Ours
Miss
Figure 10: Imputation results for our method vs. MissDiff. The imputed block is the center one. The baseline method shows characteristic white central regions with meaningful generation, con- firming our theoretical prediction (Theorem, D of learning failure in these areas.
and copy-editing suggestions, (2) formatting LaTeX code for tables and equations, (3) reviewing mathematical proofs for logical consistency and clarity, and (4) identifying and correcting typo- graphical errors throughout the manuscript.
34

---

Published as a conference paper at ICLR 2026
Imputed Data
Original Data
Imputed Data
Original Dala
(a) Mean sea level pressure (msl)
(b) Surface pressure (sp)
Imputed Data
Original Data
Imputed Data
Original Data
(c) 2-meter temperature (12m)
(d) Total column water vapor (tcwv)
Imputed Data
Original Data
Imputed Data
Original Data
(c) 100-meter U-wind component (u100)
(f) 10-meter V-wind component (v10)
Imputed Data
Original Data
Imputed Data
Original Data
(g) 100-meter V-wind component (v100)
(h) Geopotential (2)
Figure Imputation results on the ERAS dataset with 20% observed points. Each subfig- ure shows different atmospheric variable. The left column of each subfigure contains the im- puted/reconstruutte data, and the right column shows the original data.
35


Reviews:
[{"review_id": "4m8xTotACX", "reviewer": "Reviewer_ztBx", "summary": "This paper tackles a genuinely important problem: learning physical dynamics when training data is inherently incomplete. The authors propose training diffusion models by strategically splitting each incomplete sample into \"context\" (what the model sees) and \"query\" (what it tries to predict), then using ensemble averaging at inference time to reconstruct complete fields.\nThe key idea is that if you carefully design how you partition incomplete observations during training, matching the partition strategy to the underlying observation pattern (e.g., block-wise vs. pixel-wise missing data), one can train a model that learns meaningful conditional expectations even for dimensions that were never observed in the training data. The authors provide a theoretical analysis that proves convergence to the true complete data distribution and validate it on both synthetic PDEs (Shallow Water, Advection, and Navier-Stokes) and real climate data (ERA5).", "questions": "- How is K selected for ensemble averaging? Is there a principled approach to determining the number of context masks, or is selection based on empirical validation? Does the required K scale with observation sparsity?\n\n- What is the computational overhead of the proposed method? Can you provide wall-clock training and inference times relative to baseline methods? This information is essential for practical implementation.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The incomplete data setting represents a fundamental constraint in scientific applications where complete ground truth observations are physically impossible to obtain (e.g., global weather systems, ocean dynamics). The motivation is well-articulated and substantiated.\n\n- The context-query partitioning strategy constitutes a genuine advancement over existing approaches. The key differentiation with existing lies in training-time strategic partitioning that adapts to observation patterns, rather than inference-time conditioning or generic loss masking. \n\n- Theorem 1, concerning gradient scaling and parameter update frequency, provides meaningful insight into the significance of partitioning strategy selection. Theorem 2's ensemble convergence analysis demonstrates notable rigor, employing Martingale convergence theory rather than heuristic arguments. This theoretical grounding distinguishes the work from purely empirical approaches.\n\n- The performance improvements documented in Table 2 for block-wise observations are substantial. The 1% observation regime results on ERA5 are particularly noteworthy, representing an extremely challenging setting.", "weaknesses": "A bit of nitpicking:\n- The manuscript demonstrates conditions under which the method succeeds, but provides insufficient analysis of failure modes. Specifically, when does the information gap exceed the capacity of ensemble averaging to compensate? The cross-distribution experiments (Appendix G.3, Table 8) suggest performance degradation when the training and test observation ratios differ substantially; however, this warrants a more thorough discussion. Are there principled criteria for determining whether a test distribution lies outside the method's applicability range?\n\n- Ensemble averaging during inference requires K forward passes through the model (the paper employs single-step sampling with K context masks). The selection of K and the resulting computational trade-offs are not discussed. For a method targeting scientific applications, these practical considerations are important. The multi-step sampling described in Appendix E further compounds this concern. Quantitative comparison of computational costs relative to baselines would be valuable.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3A2vxmUHNJ", "reviewer": "Reviewer_VhJM", "summary": "This paper proposes a diffusion-based framework for learning physical dynamics from incomplete and irregularly sampled data, a common challenge in scientific domains. The core contribution is a novel training paradigm where each incomplete data sample is partitioned into a \"context\" set and a \"query\" set. A conditional diffusion model is then trained to reconstruct the query portions given the context, without requiring access to complete ground-truth data. The authors provide a theoretical analysis demonstrating that their method asymptotically converges to the true data distribution under certain conditions. Furthermore, they introduce an ensemble sampling technique for inference to reconstruct the complete data. The method is evaluated on several synthetic PDE datasets (Navier-Stokes, Shallow Water, Advection) and the real-world ERA5 climate dataset, showing superior performance over existing baselines, especially in highly sparse settings.", "questions": "Could the authors please elaborate on the key conceptual difference between their method and Ambient Diffusion (Daras et al., 2023)? Both seem to train diffusion models by applying masks to incomplete data. What is the crucial insight that leads to the performance gains shown?\n\nCould you provide a more concrete algorithm or pseudo-code for the \"strategic context-query partitioning\" procedure, especially for the block-wise missing data scenario? For a given observation mask M with 5/9 observed blocks, how exactly is M_ctx (e.g., with 4 blocks) sampled?\n\nRegarding the excluded baselines: Could the authors provide a more detailed justification for their exclusion? Would it be possible to conduct a qualitative or small-scale quantitative comparison to give the reader a better sense of where the proposed method stands?\n\nWhat is the inference-time cost of the ensemble sampling approach? How does the performance (e.g., MSE) and computational cost vary with the number of ensemble members K? What was the value of K used in the experiments?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper tackles the critical and practical challenge of learning physical dynamics from inherently sparse and incomplete real-world data.\n\nIt provides a theoretical analysis (Theorem 1 and 2) that offers guarantees for the learning process, explaining how and why the model can learn from incomplete data. This is a notable strength compared to more heuristic approaches.", "weaknesses": "Limited Novelty: The core idea of context-query partitioning is conceptually similar to existing masked data training strategies in self-supervised learning and other generative models. The distinction from methods like Ambient Diffusion, which also operate on incomplete data, is not sufficiently pronounced. The contribution seems to be more about a specific, effective sampling strategy rather than a fundamentally new paradigm.\n\nLack of Clarity on Implementation: The \"strategic context-query partitioning\" is the cornerstone of the method, yet its implementation is vaguely described. The paper lacks a clear algorithm or pseudo-code explaining how M_ctx is sampled from a given observation mask M for different structural patterns (e.g., block-wise). This ambiguity hinders reproducibility.\n\nIncomplete Baseline Comparison: Excluding recent methods like those by Chen et al. (2024b), Givens et al. (2025), and Zhang et al. (2025a) due to computational cost weakens the experimental validation. A more thorough comparison, even on a smaller scale, would be necessary to firmly establish state-of-the-art performance.\n\nUnaddressed Inference Cost: The proposed ensemble sampling method for reconstruction requires running the model K times per sample. This introduces a significant computational overhead at inference time, which is a critical factor for many scientific applications. The paper does not analyze this trade-off or discuss how K is chosen.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HdLkiu8OzZ", "reviewer": "Reviewer_7yXe", "summary": "This paper introduces a diffusion-based framework for learning physical dynamics when the training data itself is incomplete. The core contribution is a training strategy termed \"context-query partitioning.\" This strategy partitions the observed data within each incomplete sample into a \"context\" for model input and a \"query\" for loss computation. Theoretically, the authors demonstrate that this approach asymptotically recovers the true data distribution without full supervision, provided the partitioning strategy ensures a non-zero probability for any dimension to be queried. For inference, the paper proposes an ensemble sampling method that reconstructs the complete sample by averaging predictions over multiple, randomly drawn context masks. Empirical results on several PDE and climate datasets demonstrate superior performance over existing baselines, especially in highly sparse observation regimes.", "questions": "1. Could the authors elaborate on the conceptual similarities and differences between the \"context-query partitioning\" strategy and self-supervised learning paradigms like Masked Autoencoders (MAEs)? This is crucial for accurately positioning the paper's novelty.\n2. In the baseline comparisons, were the hyperparameters for all baseline methods tuned with the same rigor as the proposed method? Ensuring a robust comparison is key to validating the claimed performance gains.\n3. The \"distribution-preserving\" sampling strategy mentioned in line 264 appears critical to success. In real-world scenarios where prior knowledge of P_mask(M) might be inaccurate, how sensitive is the method's performance to this \"mask distribution mismatch\"?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "Problem Significance: The paper tackles a crucial and pervasive problem in scientific computing: learning from inherently sparse and incomplete observational data.\nTheoretical Rigor: The paper provides a theoretical justification (Theorem 1) for its proposed training strategy. This principled approach is a clear advantage over prior works that often rely on heuristics.\nStrong Empirical Results: The method demonstrates significant performance improvements across multiple benchmarks, particularly in sparse regimes with very low data coverage (e.g., 1%-20%).", "weaknesses": "- Insufficient Discussion of Related Work: The primary weakness is the failure to adequately discuss the connection to masked signal modeling, such as Masked Autoencoders (MAEs). The \"context-query partitioning\" is conceptually very similar to masking part of the visible data to predict another part. The absence of this discussion makes the core idea seem more novel than it is.\n- Limited Baseline Comparison: While the paper includes recent diffusion-based baselines, the modification of some baselines (e.g., adapting MissDiff to a data matching framework) makes it difficult to assess if the comparison is entirely fair. Furthermore, excluding some theoretically rigorous but computationally expensive methods, while understandable, would be more convincing with a small-scale comparison or a deeper discussion.\n- Scope of Theory: While Theorem 1 is a highlight, it focuses on asymptotic convergence. It doesn't analyze finite sample complexity or the approximation error introduced by the model architecture. The theorem's premise—that the union of query masks covers all dimensions—is satisfied by the design of the sampling strategy, which makes the argument slightly circular.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MD0rW7oOUE", "reviewer": "Reviewer_bUgQ", "summary": "This paper addressed the problem of imputing observational data under unknown dynamics. In diffusion-based imputation models, context mask and query mask are needed to train the model to extend the learned dynamics to impute unobserved data. The authors claim that the unobserved data is structured in reality and that diffusion and mask-based methods lack theoretical foundations, and then they introduce their novel pipeline to address the structured missing prior. The authors theoretically analyzed the correct strategy for context-query mask partitioning and extrapolation, and conducted comprehensive experiments on simulated and real-world datasets against competitive baseline methods. The results demonstrated the superiority of their method and the importance of catering to the structured missing prior in algorithm design. Their work could be a sound analysis platform for future diffusion-based and mask-based methods in data imputation.", "questions": "Q1: What is the term \"feature\" referring to in the main text? Does it refer to different data modalities or cell grid?\nQ2: We observed a significant improvement from block-wise to pixel-level configuration on the Navier-Stokes dataset in Table 2. Is it because the authors tuned the hyperparameters on the NS dataset, as stated in tables 6-7? This could be a source of evaluation bias.\nQ3: How is the physical dynamic learned? The data is not assumed to include the time dimension, and the \"time\" in the main text only refers to the diffusion time steps.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "S1: This paper gives a clear literature review and a contrasting comparison with similar models.\nS2: The algorithm is backed by rigorous theorems and proofs, and additionally addresses the structured missing prior in reality.\nS3: The experiments are thoughtful with ablation studies on the model design and components.", "weaknesses": "W1: The explanation of the mask selection strategy is confusing, and Figure 1 is hard to understand. In the series of training data with added noise, all data points seem to be observed once across time. No data point needs to be imputed.\nW2: Sometimes, the mask distribution $p_{mask}$ depends on the data samples $x_0$. For example, a weather station may be damaged by a typhoon's direct impact. Also, cloud blocks weather satellite observing the sea temperature, and high sea temperature also promotes cloud generation.\nW3: The authors proposed that the mask selection strategy for $M_{ctx},M_{qry}$ should align with the mask distribution prior $p_{mask}$ and used a block mask as an example. However, the authors didn't consider the case when the mask distribution pattern cannot tessellate on the whole space.\nW4: The analysis of the paper's limitations is missing in the main text.\nW5: The authors claim that some methods are not scalable, but they have not provided a scalability analysis and report on wall-time across baselines.\n\nMinor remarks:\n1. Line 723 VAE-based \"and\" GAN-based methods.\n2. The definition of $x_\\theta$ is unclear on line 138, and its input and output should be specified.\n3. What is the definition of $i$ in eq. 2? Training instance or feature dimension?\n4. A cross-reference to the proof of Thm 1 should be included in the paragraph starting from line 172.\n5. On line 361, the improvement on MissDiff from noise matching to data matching needs to be quantified.\n6. The margin of the caption of Figure 3 is too narrow.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)