Incomplete Data, Complete Dynamics: A Diffusion Approach

## Abstract

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.

## 1 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 et al., 2024; Goshisht, 2024), 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 inherently 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 is not a temporary inconvenience to be resolved through better data collection—it is an intrinsic characteristic 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 assume pixel-level independent and identically distributed (i.i.d.) missing patterns, where each spatial location has an equal probability of being observed (Daras et al., 2023; Dai et al., 2024; Simkus & Gutmann, 2025). While some recent works have explored alternative missing patterns in their experimental evaluations, such as 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 partitioning 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. Existing methods typically rely on 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-motivated approaches suffer from prohibitive computational costs, requiring multiple complete model retraining cycles or complex importance 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, we develop a theoretically principled diffusion-based framework that provides rigorous convergence guarantees while maintaining computational efficiency for high-dimensional physical dynamics problems. Our approach answers critical questions about whether diffusion models trained solely on incomplete data can recover complete data distributions, how observation patterns affect diffusion training efficiency, and under what conditions successful reconstruction of unobserved regions is guaranteed. In summary, our contributions are:

- •
Methodical design: We propose a novel conditional diffusion training paradigm that works directly with incomplete training samples, featuring a strategically designed context-query partitioning scheme tailored for physical dynamics.
- •
Theoretical guarantee: We provide the first theoretical analysis proving that diffusion-based training on incomplete data with our paradigm asymptotically recovers the true complete dynamical process under mild regularity conditions.
- •
Strong results: 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.

## 2 Preliminaries

In Appendix [A](https://arxiv.org/html/2509.20098v1#A1), 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 a detailed introduction to diffusion models in Appendix [B](https://arxiv.org/html/2509.20098v1#A2), 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 $\mathcal{X}\subset\mathbb{R}^{d}$ denote the space of complete data samples following an unknown distribution $p_{\text{data}}(\bm{x}_{0})$. Binary masks $\bm{M}\in\{0,1\}^{d}$ are drawn from distribution $p_{\text{mask}}(\bm{M})$, where 1 indicates observed elements. We assume that masks are conditionally independent of the data given the observation process: $p_{\text{mask}}(\bm{M}\mid\bm{x}_{0})=p_{\text{mask}}(\bm{M})$. In practice, we have prior knowledge about the mask distribution $p_{\text{mask}}(\bm{M})$ based on the data collection process (e.g., sensor placement patterns, measurement protocols).

For each training instance $i$, we have $\bm{x}^{(i)}_{\text{obs}}=\bm{M}^{(i)}\odot\bm{x}_{0}^{(i)}$ denoting partially observed data and $\bm{x}^{(i)}_{\text{unobs}}=(1-\bm{M}^{(i)})\odot\bm{x}_{0}^{(i)}$ representing missing values. Crucially, our training dataset $\mathcal{D}=\{(\bm{x}^{(i)}_{\text{obs}},\bm{M}^{(i)})\}_{i=1}^{N}$ contains only partial observations, no complete samples $\bm{x}_{0}^{(i)}$ are available during training. This setting reflects realistic scenarios where complete ground truth is unavailable. The objective is to learn a conditional generative model $p_{\theta}(\bm{x}_{0}\mid\bm{x}_{\text{obs}},\bm{M})$ that generates complete samples consistent with the observed elements, despite being trained solely on incomplete data.

## 3 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 to complete ground truth during training.
Our approach consists of three key components: (1) Denoising data matching on incomplete training data (Sec. [3.1](https://arxiv.org/html/2509.20098v1#S3.SS1)): We formulate a theoretically grounded training loss and establish conditions under which the model learns meaningful conditional expectations for all dimensions. (2) Strategic context-query partitioning (Sec. [3.2](https://arxiv.org/html/2509.20098v1#S3.SS2)): We develop a principled strategy for partitioning incomplete samples into context and query components, enabling reconstruction of originally missing dimensions. (3) Ensemble sampling for complete data reconstruction (Sec. [3.3](https://arxiv.org/html/2509.20098v1#S3.SS3)): We bridge the gap between training on context masks and inference on 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 on 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 a training objective that learns a conditional generative model $p_{\bm{\theta}}\left(\bm{x}_{0}\mid\bm{x}_{\text{obs}},\bm{M}\right)$(^1^11For clarity, we note that in the full diffusion imputation setting, the ideal model would output $\mathbb{E}[\bm{x}_{0}\mid\bm{x}_{t},\bm{x}_{\text{obs}},\bm{M}]$, conditioning on both the noisy state $\bm{x}_{t}$ and observations $\bm{x}_{\text{obs}}$. Our single-step sampling approach (Sec. [3.3](https://arxiv.org/html/2509.20098v1#S3.SS3)) simplifies this by using minimal noise ($t=\delta\approx 0$), effectively approximating $\mathbb{E}[\bm{x}_{0}\mid\bm{x}_{\text{obs}},\bm{M}]$. For cases requiring diversity generation, we provide a multi-step sampling procedure in Appendix [E](https://arxiv.org/html/2509.20098v1#A5) that combines both sources of information through weighted averaging.) that generates complete samples consistent 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 $(\bm{x}_{\text{obs}},\bm{M})$ in our training dataset, we treat the partially observed data $\bm{x}_{\text{obs}}$ as “complete” within the scope of available observations. We then sample context masks $\bm{M}_{\text{ctx}}\subseteq\bm{M}$ to represent “observable” portions and query masks $\bm{M}_{\text{qry}}\subseteq\bm{M}$ to represent “query” portions relative to $\bm{x}_{\text{obs}}$. Given the noisy sample $\bm{x}_{\text{obs},t}=\bm{M}\odot\left(\alpha_{t}\bm{x}_{\text{obs}}+\sigma_{t}\bm{\epsilon}\right)$ at time $t$ in the diffusion process, our training loss is formulated as:

$$ $\mathcal{L}(t,\bm{x}_{\text{obs}},\bm{M}_{\text{ctx}},\bm{M}_{\text{qry}})=\|\bm{M}_{\text{qry}}\odot(\bm{x}_{\bm{\theta}}(t,\bm{M}_{\text{ctx}}\odot\bm{x}_{\text{obs},t},\bm{M}_{\text{ctx}})-\bm{x}_{\text{obs}})\|^{2}$ (1) $$

The key architectural choice is that the model receives only the context mask $\bm{M}_{\text{ctx}}$ and corresponding observed values, trying to provide the best estimate for the queried dimensions. A natural question arises:

Since we train exclusively on incomplete data, how can this training loss enable the model to 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 $\bm{x}_{\bm{\theta}}^{*}$ be the optimal solution by minimizing the loss in equation [1](https://arxiv.org/html/2509.20098v1#S3.E1). Under the conditional independence of masks and data, we have the following results:

(i) Optimal solution: The optimal solution is given by

$$ $\left(\bm{x}_{\bm{\theta}}\left(t,\bm{M}_{\text{ctx}}\odot\bm{x}_{\text{obs},t},\bm{M}_{\text{ctx}}\right)\right)_{i}=\begin{cases}\mathbb{E}\left[\left(\bm{x}_{0}\right)_{i}\mid\bm{M}_{\text{ctx}}\odot\bm{x}_{\text{obs},t},\bm{M}_{\text{ctx}}\right],&P(\left(\bm{M}_{\text{qry}}\right)_{i}=1\mid\bm{M}_{\text{ctx}})>0\\ \text{an arbitrary value},&P(\left(\bm{M}_{\text{qry}}\right)_{i}=1\mid\bm{M}_{\text{ctx}})=0\end{cases}$ (2) $$

Specially, given the context mask $\bm{M}_{\text{ctx}}$, if the union of all possible query mask $\bm{M}_{\text{qry}}$ supports covers all feature dimensions, we have

$$ $\bm{x}_{\bm{\theta}}\left(t,\bm{M}_{\text{ctx}}\odot\bm{x}_{\text{obs},t},\bm{M}_{\text{ctx}}\right)=\mathbb{E}\left[\bm{x}_{0}\mid\bm{M}_{\text{ctx}}\odot\bm{x}_{\text{obs},t},\bm{M}_{\text{ctx}}\right]$ (3) $$

(ii) Gradient magnitude scaling: The expected squared gradient magnitude with respect to the network output for feature $i$ scales linearly with the query probability $p_{i}:=P((\bm{M}_{\text{qry}})_{i}=1\mid\bm{M}_{\text{ctx}})$:

$$ $\mathbb{E}\left[\left(\frac{\partial\mathcal{L}}{\partial(\bm{x}_{\bm{\theta}})_{i}}\right)^{2}\right]=4p_{i}\mathbb{E}\left[\left((\bm{x}_{\bm{\theta}})_{i}-(\bm{x}_{\text{obs}})_{i}\right)^{2}\Big|(\bm{M}_{\text{qry}})_{i}=1\right]$ (4) $$

(iii) Parameter update frequency: The frequency of non-zero parameter updates for feature $i$ is exactly $p_{i}$:

$$ $P(\text{feature }i\text{ contributes to parameter update})=P((\bm{M}_{\text{qry}})_{i}=1\mid\bm{M}_{\text{ctx}})=p_{i}$ (5) $$

This theorem reveals a critical insight: the model learns meaningful conditional expectations $\mathbb{E}\left[\left(\bm{x}_{0}\right)_{i}\mid\bm{M}_{\text{ctx}}\odot\bm{x}_{\text{obs},t},\bm{M}_{\text{ctx}}\right]$ for dimension $i$ only when $P(\left(\bm{M}_{\text{qry}}\right)_{i}=1\mid\bm{M}_{\text{ctx}})>0$. When this probability is zero, the model’s output for dimension $i$ becomes arbitrary since it never receives gradient updates for that dimension. This theoretical finding directly drives the need for strategic context-query partitioning: given any context mask $\bm{M}_{\text{ctx}}$ (without information of $\bm{M}$), we must ensure that every dimension outside the context, including dimensions that were originally missing in the raw data (i.e., dimensions $i$ where $\bm{M}_{i}=0$), has a positive probability of being selected as a query point.
In the following section, we detail how to design the context-query mask sampling strategy to guarantee this requirement.

###### Theorem 1 (Optimal solution under context masking without query information) .

### 3.2 Strategic context-query 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 a principled design framework that guarantees positive query probabilities for all observable dimensions.

#### Design principle.

Based on Theorem [1](https://arxiv.org/html/2509.20098v1#Thmtheorem1), we establish the core design principle for effective context-query partitioning:

###### Principle 1 (Principle of uniform query exposure) .

For effective learning from incomplete data, the context-query partitioning strategy must satisfy:

- 1.
Non-zero query probability: For all unobserved dimensions $i$, i.e., $\left(\bm{M}_{\text{ctx}}\right)_{i}=0$,
$P((\bm{M}_{\text{qry}})_{i}=1\mid\bm{M}_{\text{ctx}})>0$
(6)
- 2.
Uniform exposure: The query probabilities should be approximately uniform across all observed dimensions to achieve balanced learning:
$P((\bm{M}_{\text{qry}})_{i}=1\mid\bm{M}_{\text{ctx}})\approx P((\bm{M}_{\text{qry}})_{j}=1\mid\bm{M}_{\text{ctx}})\quad\forall i,j:\left(\bm{M}_{\text{ctx}}\right)_{i}=\left(\bm{M}_{\text{ctx}}\right)_{j}=0$
(7)

To implement this principle, we can decompose the query probability using the law of total probability over all possible observation masks:

$$ $P((\bm{M}_{\text{qry}})_{i}=1\mid\bm{M}_{\text{ctx}})=\sum_{\bm{M}}P((\bm{M}_{\text{qry}})_{i}=1\mid\bm{M}_{\text{ctx}},\bm{M})\cdot P(\bm{M}\mid\bm{M}_{\text{ctx}})$ (8) $$

This decomposition reveals that the query probability depends on two factors: (1) the conditional query sampling strategy given both context and observation masks $P((\bm{M}_{\text{qry}})_{i}=1\mid\bm{M}_{\text{ctx}},\bm{M})$, and (2) the posterior distribution of observation masks given the context $P(\bm{M}\mid\bm{M}_{\text{ctx}})$.
Since $\bm{M}_{\text{ctx}}\subseteq\bm{M}$ by construction, we can always find observation masks $\bm{M}$ such that $P(\bm{M}\mid\bm{M}_{\text{ctx}})>0$. However, $P((\bm{M}_{\text{qry}})_{i}=1\mid\bm{M}_{\text{ctx}},\bm{M})$ is not guaranteed to be positive for all dimensions $i$. Fortunately, we can strategically design the sampling mechanism for $\bm{M}_{\text{ctx}}$ to ensure that there indeed exist observation masks $\bm{M}$ where both terms are simultaneously positive, thereby guaranteeing $P((\bm{M}_{\text{qry}})_{i}=1\mid\bm{M}_{\text{ctx}})>0$ for all observed dimensions.
To illustrate how different context mask selection strategies affect the query probability $P((\bm{M}_{\text{qry}})_{i}=1\mid\bm{M}_{\text{ctx}})$, we first present a concrete example that demonstrates the critical impact of this design choice.

Figure: 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-structured sampling ensures balanced learning across all dimensions. See Fig. [10](https://arxiv.org/html/2509.20098v1#A7.F10) for the resulting imputation failures.
Refer to caption: x1.png

Illustrate example. Fig. [1](https://arxiv.org/html/2509.20098v1#S3.F1) demonstrates how different partitioning strategies affect learning effectiveness using block-structured observation patterns.
Consider a scenario where observation masks $\bm{M}$ randomly mask 2 out of 9 spatial blocks:

- •
Problematic strategy (top): For each partially observed sample $\bm{x}_{\text{obs}}$, context points are selected by uniform sampling across all observable dimensions, yielding $\bm{M}_{\text{ctx}}^{\text{uni}}$. This strategy is problematic because for such given $\bm{M}_{\text{ctx}}^{\text{uni}}$, there exists only one observation mask $\bm{M}$ that contains it (i.e., the specific $\bm{M}$ from which $\bm{x}_{\text{obs}}$ was derived). Consequently, for masked dimensions $i$ where $\bm{M}_{i}=0$, we have $P((\bm{M}_{\text{qry}})_{i}=1\mid\bm{M}_{\text{ctx}}^{\text{uni}},\bm{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 $\bm{M}_{\text{ctx}}^{\text{block}}$ that typically contains 4 complete blocks. Crucially, for a given $\bm{M}_{\text{ctx}}^{\text{block}}$, there exist multiple possible observation masks $\bm{M}$ that contain it (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 $\bm{M}$ such that $P((\bm{M}_{\text{qry}})_{i}=1\mid\bm{M}_{\text{ctx}}^{\text{block}},\bm{M})>0$, thereby guaranteeing positive query probabilities across all observable dimensions.

###### Principle 1 (Principle of uniform query exposure) .

#### 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 or reasonable prior knowledge about $p_{\text{mask}}(\bm{M})$. In practice, to satisfy Principle [1](https://arxiv.org/html/2509.20098v1#Thmprinciple1), a viable strategy is to design the context mask sampling mechanism $\bm{M}_{\text{ctx}}$ based on the observation mask distribution $p_{\text{mask}}(\bm{M})$. Specifically, we sample $\bm{M}_{\text{ctx}}$ from $\bm{M}$ following the same structural pattern as $p_{\text{mask}}(\bm{M})$: for i.i.d. pixel-level observations, we independently sample each observed pixel; for block-structured observations, we sample complete blocks from available blocks in $\bm{M}$. This distribution-preserving strategy ensures that every observed dimension can potentially be excluded from context (and included in query), guaranteeing $P\left(\left(\bm{M}_{\text{qry}}\right)_{i}=1\mid\bm{M}_{\text{ctx}}\right)>0$ for all $i$. Under this design paradigm, the model learns during training to recover $\bm{x}_{0}$ from context observations while leveraging the structural knowledge embedded in $p_{\text{mask}}(\bm{M})$. At inference time, the model can then utilize the complete observation $\bm{x}_{\text{obs}}$ along with the same distributional knowledge $p_{\text{mask}}(\bm{M})$ to reconstruct the full data $\bm{x}_{0}$.

Figure: Algorithm 1 Diffusion-based training for missing data imputation

#### Training algorithm and practical considerations.

Building on our theoretical foundation and design principles, we present our training algorithm in Alg. [1](https://arxiv.org/html/2509.20098v1#alg1). While ensuring $P((\bm{M}_{\text{qry}})_{i}=1\mid\bm{M}_{\text{ctx}})>0$ is necessary, the choice of context-query ratio involves two critical trade-offs:

- •
Information gap trade-off (Theorem [2](https://arxiv.org/html/2509.20098v1#Thmtheorem2)): When $\bm{M}_{\text{ctx}}$ contains few observed points relative to $\bm{M}$, the large information gap increases approximation variance and slows down convergence.
- •
Parameter update frequency trade-off (equation [5](https://arxiv.org/html/2509.20098v1#S3.E5)): When $\bm{M}_{\text{ctx}}$ contains too many observed points, query probabilities $p_{i}$ become small, leading to infrequent parameter updates for reconstructing missing information.

These theoretical considerations suggest that moderate context ratios should achieve optimal performance by balancing both trade-offs. Our experimental results confirm this theoretical prediction, with detailed analysis provided in Appendix [G.2](https://arxiv.org/html/2509.20098v1#A7.SS2) and Tab. [6](https://arxiv.org/html/2509.20098v1#A7.T6).

### 3.3 Ensemble sampling for complete data reconstruction

Our trained model approximates the conditional expectation $\mathbb{E}[\bm{x}_{0}\mid\bm{M}_{\text{ctx}}\odot\bm{x}_{\text{obs},t},\bm{M}_{\text{ctx}}]$ given a randomly sampled context mask, rather than the desired full conditional expectation $\mathbb{E}[\bm{x}_{0}\mid\bm{x}_{\text{obs},t},\bm{M}]$ that conditions on the complete observation. To bridge this gap and enable complete data reconstruction, 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 informative to constrain the solutions to a relatively concentrated region in the solution space (Alberti & Santacesaria, 2021).
When the posterior distribution $p(\bm{x}_{0}\mid\bm{x}_{\text{obs}},\bm{M})$ is highly concentrated with solutions clustered closely together, it can be well-approximated by a narrow distribution centered at $\bm{x}^{*}$, where $\bm{x}^{*}$ represents the mean of the tightly clustered solutions consistent with the observations. In such cases, $\mathbb{E}[\bm{x}_{0}\mid\bm{x}_{\text{obs}},\bm{M}]\approx\bm{x}^{*}$ provides a good representative solution, reducing the need for extensive iterative denoising steps.

To leverage this property, we implement a single-step sampling procedure. We apply minimal noise at timestep $t=\delta$ where $0<\delta\ll 1$:

$$ $\bm{x}_{\delta}=\alpha_{\delta}\bm{x}_{\text{obs}}+\sigma_{\delta}\bm{\epsilon},\quad\bm{\epsilon}\sim\mathcal{N}(\bm{0},\bm{I})$ (9) $$

The small noise level ensures $\bm{M}\odot\bm{x}_{\delta}\approx\bm{x}_{\text{obs}}$. We then approximate the desired conditional expectation using ensemble averaging over $K$ randomly sampled context masks:

$$ $\bm{x}^{*}=\mathbb{E}[\bm{x}_{0}\mid\bm{x}_{\text{obs}},\bm{M}]\approx\frac{1}{K}\sum_{k=1}^{K}\bm{x}_{\theta}\left(\delta,\bm{M}_{\text{ctx}}^{(k)}\odot\bm{x}_{\text{obs},\delta},\bm{M}_{\text{ctx}}^{(k)}\right),$ (10) $$

where $\bm{M}_{\text{ctx}}^{(1)},\ldots,\bm{M}_{\text{ctx}}^{(K)}\subseteq\bm{M}$ are conditionally independent given $\bm{M}$.
This approach enables direct reconstruction in a single denoising step, making it particularly suitable for well-posed inverse problems where the observations strongly constrain the solution space.

We then establish the theoretical foundation that justifies our ensemble averaging approach. Let $\text{obs}=[\bm{M}\odot\bm{x}_{\text{obs},t},\bm{M}]$ and $\text{ctx}=[\bm{M}_{\text{ctx}}\odot\bm{x}_{\text{obs},t},\bm{M}_{\text{ctx}}]$ denote the full observation and context observation respectively.

###### Theorem 2 (Ensemble approximation convergence) .

Let $\mathbb{E}[\bm{x}_{0}\mid\mathrm{obs}]:=\mathbb{E}[\bm{x}_{0}\mid\bm{x}_{\mathrm{obs},t},\bm{M}]$ be the ground truth conditional expectation and $\mathbb{E}[\bm{x}_{0}\mid\mathrm{ctx}]=\mathbb{E}[\bm{x}_{0}\mid\bm{M}_{\text{ctx}}\odot\bm{x}_{\text{obs},t},\bm{M}_{\text{ctx}}]$ be the context-conditioned expectation. The expected squared error between these quantities is:

$$ $\mathbb{E}\left[\|\mathbb{E}[\bm{x}_{0}\mid\mathrm{ctx}]-\mathbb{E}[\bm{x}_{0}\mid\mathrm{obs}]\|^{2}\right]=\mathbb{E}[\mathrm{Var}[\bm{x}_{0}\mid\mathrm{ctx}]]-\mathbb{E}[\mathrm{Var}[\bm{x}_{0}\mid\mathrm{obs}]]$ (11) $$

Consider a practical model with output $\bm{x}_{\bm{\theta}}(t,\mathrm{ctx})=\mathbb{E}[\bm{x}_{0}\mid\mathrm{ctx}]+\bm{b}(\mathrm{ctx})+\bm{\epsilon}_{\text{bias}}(\mathrm{ctx})$, where $\bm{b}(\mathrm{ctx})$ represents context-dependent deterministic bias and $\bm{\epsilon}_{\text{bias}}(\mathrm{ctx})$ is random error with $\mathbb{E}[\bm{\epsilon}_{\text{bias}}]=\bm{0}$. Given the ensemble prediction in equation [10](https://arxiv.org/html/2509.20098v1#S3.E10) as:

$$ $\hat{\bm{\mu}}_{K}=\frac{1}{K}\sum_{k=1}^{K}\bm{x}_{\bm{\theta}}(t,\mathrm{ctx}^{(k)}),$ (12) $$

the expected squared error between the ensemble prediction and ground truth is:

$$ $\begin{split}\mathbb{E}\left[\left\|\hat{\bm{\mu}}_{K}-\mathbb{E}[\bm{x}_{0}\mid\mathrm{obs}]\right\|^{2}\right]&=\|\underbrace{\mathbb{E}[\mathbb{E}[\bm{x}_{0}\mid\mathrm{ctx}]]-\mathbb{E}[\bm{x}_{0}\mid\mathrm{obs}]}_{\text{information gap}}+\underbrace{\mathbb{E}[\bm{b}(\mathrm{ctx})]}_{\text{model bias}}\|^{2}\\ &\quad+\frac{1}{K}(\underbrace{\mathrm{Var}[\mathbb{E}[\bm{x}_{0}\mid\mathrm{ctx}]]}_{\text{data variance}}+\underbrace{\mathrm{Var}[\bm{b}(\mathrm{ctx})]+\mathrm{Var}[\bm{\epsilon}_{\text{bias}}]}_{\text{model variance}})\end{split}$ (13) $$

As $K\to\infty$, the ensemble converges to:

$$ $\lim_{K\to\infty}\mathbb{E}\left[\left\|\hat{\bm{\mu}}_{K}-\mathbb{E}[\bm{x}_{0}\mid\mathrm{obs}]\right\|^{2}\right]=\left\|\mathbb{E}[\mathbb{E}[\bm{x}_{0}\mid\mathrm{ctx}]]-\mathbb{E}[\bm{x}_{0}\mid\mathrm{obs}]+\mathbb{E}[\bm{b}(\mathrm{ctx})]\right\|^{2}$ (14) $$

This theorem demonstrates that ensemble averaging eliminates the variance terms, with the remaining error determined by the information gap between context and full observations, plus any systematic model bias. The proof is provided in Appendix [F.2](https://arxiv.org/html/2509.20098v1#A6.SS2).

###### Theorem 2 (Ensemble approximation convergence) .

#### Multi-step sampling.

While single-step sampling suffices when observations nearly determine a unique solution, generating diverse imputed samples or 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 [E](https://arxiv.org/html/2509.20098v1#A5).

## 4 Experiment

### 4.1 Baselines

We compare against established baselines, including traditional imputation methods (Temporal Consistency (Huang et al., 2016), Fast Marching (Telea, 2004), Navier-Stokes inpainting (Bertalmio et al., 2001)) and recent diffusion-based approaches, MissDiff (Ouyang et al., 2023), AmbientDiff (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. This modification makes MissDiff more comparable to AmbientDiff, since both methods now use the same data matching framework. As a result, all three diffusion-based approaches, MissDiff, AmbientDiff, and our method, now employ the data matching paradigm. For all 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 low-dimensional data and incur prohibitively high training costs when scaled to the high-dimensional datasets used in our experiments.

### 4.2 Datasets

We conduct comprehensive experiments to validate our theoretical framework and demonstrate the effectiveness of 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., 2020), under varying levels of data sparsity, ranging from 80% to as low as 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 is never available during training, distinguishing our setting from conventional imputation tasks that artificially mask complete observations.

Figure: (a) Visualization of data imputation results on Shallow Water and Advection datasets. The training set contains 30% observed data points. Each example shows a complete sample reconstructed by the model based on partial observations. Observed entries are marked with rectangles, while missing entries are filled in by the model.
Refer to caption: figs/sw/sw_5frames_0.png

Shallow Water and Advection Equations. We consider two fundamental geophysical PDE systems: 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 1k test samples with $32\times 32$ spatial resolution and 50 temporal frames, generated with randomized physical parameters and initial conditions.

Navier-Stokes Equations. We use incompressible Navier-Stokes simulations of isotropic turbulence, featuring the characteristic Kolmogorov energy cascade. The dataset comprises 1,152 samples at $64\times 64$ resolution with 100-frame sequences, generated using spectral or finite volume methods.

ERA5 Reanalysis. For real-world evaluation, we utilize the ERA5 atmospheric reanalysis from ECMWF, incorporating nine essential meteorological variables. We process one year of hourly data at $103\times 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. Detailed dataset specifications and evaluation protocols are provided in Appendix [G.1](https://arxiv.org/html/2509.20098v1#A7.SS1).

Figure: 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 at $64\times 64$ resolution.
Refer to caption: x3.png

**Table 1: 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.**
| Method | Navier-Stokes ($\times 10^{-3}$) | ERA5 ($\times 10^{-2}$) |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| 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 $\pm$ 0.025 | 0.611 $\pm$ 0.077 | 3.077 $\pm$ 1.046 | 0.416 $\pm$ 0.004 | 0.676 $\pm$ 0.088 | 1.653 $\pm$ 0.296 |
| AmbientDiff | 0.238 $\pm$ 0.017 | 0.538 $\pm$ 0.024 | 2.043 $\pm$ 0.089 | 0.256 $\pm$ 0.002 | 0.414 $\pm$ 0.031 | 1.234 $\pm$ 0.437 |
| Ours | 0.223 $\pm$ 0.016 | 0.507 $\pm$ 0.026 | 1.931 $\pm$ 0.092 | 0.250 $\pm$ 0.002 | 0.408 $\pm$ 0.030 | 1.229 $\pm$ 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-wise context-query partitioning strategies.**
| Method | Shallow Water | Advection | Navier-Stokes |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| 8/9 | 5/9 | 8/9 | 5/9 | 8/9 |  |  |
| Temporal Consistency | 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 $\pm$ 0.0024 | 0.1166 $\pm$ 0.0066 | 0.1202 $\pm$ 0.0047 | 0.1979 $\pm$ 0.0228 | 1.4357 $\pm$ 0.1132 |  |
| AmbientDiff | 0.0217 $\pm$ 0.0063 | 0.0925 $\pm$ 0.0017 | 0.1077 $\pm$ 0.0009 | 0.1524 $\pm$ 0.0137 | 1.4954 $\pm$ 0.2609 |  |
| Ours | pixel-level (incorrect) | 0.0215 $\pm$ 0.0035 | 0.0989 $\pm$ 0.0007 | 0.1171 $\pm$ 0.0041 | 0.1894 $\pm$ 0.0179 | 1.4925 $\pm$ 0.2609 |
| block-wise (correct) | 0.0203 $\pm$ 0.0059 | 0.0865 $\pm$ 0.0014 | 0.1065 $\pm$ 0.0009 | 0.1407 $\pm$ 0.0116 | 0.7592 $\pm$ 0.0386 |  |

### 4.3 Experiment results

#### Pixel-level observation.

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 correlations. Tab. [1](https://arxiv.org/html/2509.20098v1#S4.T1) and Tab. [9](https://arxiv.org/html/2509.20098v1#A7.T9) demonstrate that our method achieves superior performance in the majority of evaluation scenarios.

#### 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](https://arxiv.org/html/2509.20098v1#S4.T2), demonstrate that the proposed strategic context-query partitioning, which adapts to the observation pattern during training, is essential. When the partitioning strategy matches the observation structure (block masks), our method effectively learns to reconstruct complete fields. Conversely, mismatched strategies lead to degraded performance, validating our theoretical analysis.

#### Cross-distribution generalization.

Beyond evaluating on matching train-test distributions, we also investigate cross-distribution generalization where models trained with one observation ratio (e.g., 80%) are tested on data with different observation densities (e.g., 60% or 20%). We provide detailed experiment results and analysis of this cross-distribution setting and our adaptive context sampling strategy in Appendix [G.3](https://arxiv.org/html/2509.20098v1#A7.SS3).

### 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 $\bm{M}_{\text{ctx}}$ with $\bm{M}$, (2) the choice of backbone architecture, and (3) the context and query mask ratio selection guided by our theoretical analysis. Detailed results and analysis are provided in Appendix [G.2](https://arxiv.org/html/2509.20098v1#A7.SS2). Through systematic ablation studies, we validate our key theoretical insights. The results demonstrate that our method consistently outperforms existing approaches while providing theoretical guarantees for convergence to the desired conditional expectations.

## 5 Conclusion

We presented a principled framework for learning physical dynamics from incomplete observations using diffusion models trained directly on partial data through strategic context-query partitioning. Our theoretical analysis proves that training on incomplete 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 effectiveness across diverse physical systems demonstrates practical applicability for real-world scenarios where complete observations are inherently unavailable.

## 6 Reproducibility statement

To ensure the reproducibility of our work, we will make all experimental code and datasets publicly available upon publication. The complete implementation, including model architectures, training procedures, and evaluation scripts, will be provided as supplementary materials with detailed documentation. All datasets used in our experiments will be released, with the exception of the ERA5 dataset, which requires individual registration and download from the official ECMWF website due to licensing restrictions.

## References

- Albergo et al. (2023)
Michael S Albergo, Mark Goldstein, Nicholas M Boffi, Rajesh Ranganath, and Eric Vanden-Eijnden.
Stochastic interpolants with data-dependent couplings.
*arXiv preprint arXiv:2310.03725*, 2023.
- Alberti & Santacesaria (2021)
Giovanni S Alberti and Matteo Santacesaria.
Infinite dimensional compressed sensing from anisotropic measurements and applications to inverse problems in pde.
*Applied and Computational Harmonic Analysis*, 50:105–146, 2021.
- Bertalmio et al. (2001)
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.
- Bishop & Nasrabadi (2006)
Christopher M Bishop and Nasser M Nasrabadi.
*Pattern recognition and machine learning*, volume 4.
Springer, 2006.
- Brunton & Kutz (2024)
Steven L Brunton and J Nathan Kutz.
Promising directions of machine learning for partial differential equations.
*Nature Computational Science*, 4(7):483–494, 2024.
- Cao (2024)
Shuhao Cao.
Navier-stokes dataset of isotropic turbulence in a periodic box, 2024.
URL [https://huggingface.co/datasets/scaomath/navier-stokes-dataset](https://huggingface.co/datasets/scaomath/navier-stokes-dataset).
Funded by National Science Foundation: NSF award DMS-2309778.
- Chen et al. (2024a)
Hanyang Chen, Yang Jiang, Shengnan Guo, Xiaowei Mao, Youfang Lin, and Huaiyu Wan.
Difflight: a partial rewards conditioned diffusion model for traffic signal control with missing data.
*Advances in Neural Information Processing Systems*, 37:123353–123378, 2024a.
- Chen et al. (2024b)
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 in Neural Information Processing Systems*, 37:112050–112103, 2024b.
- Chung et al. (2022)
Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Marc L Klasky, and Jong Chul Ye.
Diffusion posterior sampling for general noisy inverse problems.
*arXiv preprint arXiv:2209.14687*, 2022.
- Conti (2024)
Silvia Conti.
Artificial intelligence for weather forecasting.
*Nature Reviews Electrical Engineering*, 1(1):8–8, 2024.
- Corneanu et al. (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.
- Dai et al. (2024)
Zongyu Dai, Emily Getzen, and Qi Long.
Sadi: Similarity-aware diffusion model-based imputation for incomplete temporal ehr data.
In *International Conference on Artificial Intelligence and Statistics*, pp. 4195–4203. PMLR, 2024.
- Daras et al. (2023)
Giannis Daras, Kulin Shah, Yuval Dagan, Aravind Gollakota, Alex Dimakis, and Adam Klivans.
Ambient diffusion: Learning clean distributions from corrupted data.
*Advances in Neural Information Processing Systems*, 36:288–313, 2023.
- Du et al. (2021)
Yilun Du, Katie Collins, Josh Tenenbaum, and Vincent Sitzmann.
Learning signal-agnostic manifolds of neural fields.
*Advances in Neural Information Processing Systems*, 34:8320–8331, 2021.
- Duan et al. (2024)
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.
- Dupont et al. (2022)
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 it like one.
*arXiv preprint arXiv:2201.12204*, 2022.
- Efron (2011)
Bradley Efron.
Tweedie’s formula and selection bias.
*Journal of the American Statistical Association*, 106(496):1602–1614, 2011.
- Givens et al. (2025)
Josh Givens, Song Liu, and Henry WJ Reeve.
Score matching with missing data.
*arXiv preprint arXiv:2506.00557*, 2025.
- Goshisht (2024)
Manoj Kumar Goshisht.
Machine learning and deep learning in synthetic biology: Key architectures, applications, and challenges.
*ACS omega*, 9(9):9921–9945, 2024.
- Hersbach et al. (2020)
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*, 146(730):1999–2049, 2020.
- Ho et al. (2020)
Jonathan Ho, Ajay Jain, and Pieter Abbeel.
Denoising diffusion probabilistic models.
*Advances in neural information processing systems*, 2020.
- Huang et al. (2016)
Jia-Bin Huang, Sing Bing Kang, Narendra Ahuja, and Johannes Kopf.
Temporally coherent completion of dynamic video.
*ACM Transactions on Graphics (ToG)*, 35(6):1–11, 2016.
- Huang et al. (2024)
Jiahe Huang, Guandao Yang, Zichen Wang, and Jeong Joon Park.
Diffusionpde: Generative pde-solving under partial observation.
*arXiv preprint arXiv:2406.17763*, 2024.
- Ipsen et al. (2020)
Niels Bruun Ipsen, Pierre-Alexandre Mattei, and Jes Frellsen.
not-miwae: Deep generative modelling with missing not at random data.
*arXiv preprint arXiv:2006.12871*, 2020.
- Karras et al. (2024)
Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine.
Analyzing and improving the training dynamics of diffusion models.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pp. 24174–24184, 2024.
- Kim & Ye (2021)
Kwanyoung Kim and Jong Chul Ye.
Noise2score: tweedie’s approach to self-supervised image denoising without clean images.
*Advances in Neural Information Processing Systems*, 34:864–874, 2021.
- Klöwer et al. (2018)
M. Klöwer, M. F. Jansen, M. Claus, R. J. Greatbatch, and S. Thomsen.
Energy budget-based backscatter in a shallow water model of a double gyre basin.
*Ocean Modelling*, 132, 2018.
doi: 10.1016/j.ocemod.2018.09.006.
- Lemercier et al. (2025)
Jean-Marie Lemercier, Julius Richter, Simon Welker, Eloi Moliner, Vesa Välimäki, and Timo Gerkmann.
Diffusion models for audio restoration: A review [special issue on model-based and data-driven audio signal processing].
*IEEE Signal Processing Magazine*, 41(6):72–84, 2025.
- Li & Marlin (2020)
Steven Cheng-Xian Li and Benjamin Marlin.
Learning from irregularly-sampled time series: A missing data perspective.
In *International Conference on Machine Learning*, pp. 5937–5946. PMLR, 2020.
- Li et al. (2019)
Steven Cheng-Xian Li, Bo Jiang, and Benjamin Marlin.
Misgan: Learning from incomplete data with generative adversarial networks.
*arXiv preprint arXiv:1902.09599*, 2019.
- Li et al. (2020)
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar.
Fourier neural operator for parametric partial differential equations.
*arXiv preprint arXiv:2010.08895*, 2020.
- Liu et al. (2023)
Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos A Theodorou, Weili Nie, and Anima Anandkumar.
I^2sb: Image-to-image schrödinger bridge.
*arXiv preprint arXiv:2302.05872*, 2023.
- Lugmayr et al. (2022)
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.
- Luo et al. (2025)
Kuang Luo, Jingshang Zhao, Yingping Wang, Jiayao Li, Junjie Wen, Jiong Liang, Henry Soekmadji, and Shaolin Liao.
Physics-informed neural networks for pde problems: a comprehensive review.
*Artificial Intelligence Review*, 58(10):1–43, 2025.
- Ma et al. (2020)
Chao Ma, Sebastian Tschiatschek, Richard Turner, José Miguel Hernández-Lobato, and Cheng Zhang.
Vaem: a deep generative model for heterogeneous mixed type data.
*Advances in Neural Information Processing Systems*, 33:11237–11247, 2020.
- Meng et al. (2025)
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 Computational Science and Engineering*, 1(1):20, 2025.
- Ouyang et al. (2023)
Yidong Ouyang, Liyan Xie, Chongxuan Li, and Guang Cheng.
Missdiff: Training diffusion models on tabular data with missing values.
*arXiv preprint arXiv:2307.00467*, 2023.
- Qi et al. (2024)
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.
- Saharia et al. (2022)
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.
- Simkus & Gutmann (2025)
Vaidotas Simkus and Michael U Gutmann.
Cfmi: Flow matching for missing data imputation.
*arXiv preprint arXiv:2506.09258*, 2025.
- Song & Ermon (2019)
Yang Song and Stefano Ermon.
Generative modeling by estimating gradients of the data distribution.
*Advances in neural information processing systems*, 32, 2019.
- Song et al. (2020)
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.
- Telea (2004)
Alexandru Telea.
An image inpainting technique based on the fast marching method.
*Journal of graphics tools*, 9(1):23–34, 2004.
- Wang et al. (2024)
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.12171*, 2024.
- Xu et al. (2023)
Yilun Xu, Mingyang Deng, Xiang Cheng, Yonglong Tian, Ziming Liu, and Tommi Jaakkola.
Restart sampling for improving generative processes.
*Advances in Neural Information Processing Systems*, 36:76806–76838, 2023.
- Xue et al. (2024)
Kaiwen Xue, Yuhao Zhou, Shen Nie, Xu Min, Xiaolu Zhang, Jun Zhou, and Chongxuan Li.
Unifying bayesian flow networks and diffusion models through stochastic differential equations.
*arXiv preprint arXiv:2404.15766*, 2024.
- Yue et al. (2023)
Conghan Yue, Zhengwei Peng, Junlong Ma, Shiyan Du, Pengxu Wei, and Dongyu Zhang.
Image restoration through generalized ornstein-uhlenbeck bridge.
*arXiv preprint arXiv:2312.10299*, 2023.
- Zhang et al. (2025a)
Hengrui Zhang, Liancheng Fang, Qitian Wu, and Philip S Yu.
Diffputer: Empowering diffusion models for missing data imputation.
In *The Thirteenth International Conference on Learning Representations*, 2025a.
- Zhang et al. (2025b)
Huijun Zhang, Yaxin Liu, Chongyu Zhang, and Ningyun Li.
Machine learning methods for weather forecasting: A survey.
*Atmosphere*, 16(1):82, 2025b.
- Zheng et al. (2023)
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.
- Zhou et al. (2024)
Zihan Zhou, Xiaoxue Wang, and Tianshu Yu.
Generating physical dynamics under priors.
*arXiv preprint arXiv:2409.00730*, 2024.
- Zhuang et al. (2023)
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.

## Appendix A Related work

### A.1 Imputation

Imputation, the task of filling missing or corrupted regions in data with plausible content, has been revolutionized by diffusion models across various modalities (Corneanu et al., 2024; Lemercier et al., 2025; Duan et al., 2024). Current imputation approaches follow three primary paradigms. Palette (Saharia 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-free method that leverages pretrained unconditional models by resampling known regions while generating content only for masked areas. Bridge-based methods (Liu et al., 2023; Yue et al., 2023; Albergo et al., 2023) design specialized diffusion processes between original and masked data distributions, requiring models trained to condition directly on masked inputs. DiffusionPDE (Huang et al., 2024) introduces a diffusion model to solve PDEs under partial observation by learning the joint distribution 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), 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., 2024a) 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) proposed 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 real-world scenarios.

## Appendix B Data matching diffusion models

Diffusion models (Song & 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 $\bm{x}_{0}\sim p_{0}$ according to the stochastic differential equation:

$$ $\mathrm{d}\bm{x}_{t}=f(t)\bm{x}_{t}\mathrm{d}t+g(t)\mathrm{d}\bm{w}_{t},\quad\bm{x}_{0}\sim p_{0},$ (15) $$

where $\bm{w}_{t}\in\mathbb{R}^{d}$ represents standard Brownian motion, $f(t)=\frac{\mathrm{d}\log\alpha_{t}}{\mathrm{d}t}$, $g^{2}(t)=\frac{\mathrm{d}\sigma_{t}^{2}}{\mathrm{d}t}-2\frac{\mathrm{d}\log\alpha_{t}}{\mathrm{d}t}\sigma_{t}^{2}$, and $\alpha_{t}$, $\sigma_{t}$ are predefined time-dependent functions. This process has the analytical solution $p_{t}\left(\bm{x}_{t}\mid\bm{x}_{0}\right)=\mathcal{N}\left(\bm{x}_{t}\mid\alpha_{t}\bm{x}_{0},\sigma_{t}^{2}\bm{I}\right)$.

The reverse process generates samples by integrating backward from noise to data using:

$$ $\frac{\mathrm{d}\bm{x}_{t}}{\mathrm{d}t}=f(t)\bm{x}_{t}-\frac{1}{2}g^{2}(t)\nabla_{\bm{x}}\log p_{t}(\bm{x}_{t}),\quad\bm{x}_{T}\sim p_{T}(\bm{x}_{T}).$ (16) $$

Since the terminal distribution $p_{T}(\bm{x}_{T})$ becomes approximately Gaussian through appropriate parameter choices, sampling from it and reversing the process yields samples from $p_{0}$. The score function $\nabla_{\bm{x}}\log p_{t}(\bm{x}_{t})$ is computationally intractable but can be estimated using neural networks via noise matching and data matching approaches (Zheng et al., 2023):

$$ $\displaystyle\mathcal{J}_{\textsubscript{noise}}(\bm{\theta})$ $\displaystyle=\mathbb{E}_{t,\bm{x}_{0},\bm{\epsilon}}\left[w(t)\left\|\bm{\epsilon}_{\bm{\theta}}\left(\bm{x}_{t},t\right)-\bm{\epsilon}\right\|^{2}\right],\quad$ $\displaystyle\bm{\epsilon}_{\bm{\theta}}^{*}\left(\bm{x}_{t},t\right)$ $\displaystyle=-\sigma_{t}\nabla_{\bm{x}}\log p_{t}\left(\bm{x}_{t}\right);$ (17a) $\displaystyle\mathcal{J}_{\textsubscript{data}}(\bm{\theta})$ $\displaystyle=\mathbb{E}_{t,\bm{x}_{0},\bm{\epsilon}}\left[w(t)\left\|\bm{x}_{\bm{\theta}}\left(\bm{x}_{t},t\right)-\bm{x}_{0}\right\|^{2}\right],\quad$ $\displaystyle\bm{x}_{\bm{\theta}}^{*}\left(\bm{x}_{t},t\right)$ $\displaystyle=\frac{1}{\alpha_{t}}\bm{x}_{t}+\frac{\sigma_{t}^{2}}{\alpha_{t}}\nabla_{\bm{x}}\log p_{t}\left(\bm{x}_{t}\right),$ (17b) $$

where $\bm{x}_{t}=\alpha_{t}\bm{x}_{0}+\sigma_{t}\bm{\epsilon}$ with $\bm{\epsilon}\sim\mathcal{N}\left(\bm{0},\bm{I}\right)$, and $w(t)$ is a loss weight function. By Tweedie’s formula (Efron, 2011; Kim & Ye, 2021; Chung et al., 2022), we also have $\bm{x}_{\bm{\theta}}^{*}\left(\bm{x}_{t},t\right)=\mathbb{E}[\bm{x}_{0}\mid\bm{x}_{t}]$.

## Appendix C Method comparison and theoretical analysis

### C.1 Diffusion probabilistic fields

Several prior studies have proposed training generative models using field representations (Du et al., 2021; Dupont et al., 2022; Zhuang et al., 2023). Similarly, our approach trains a model to predict $\mathbb{E}\left[\bm{x}_{0}\mid\bm{x}_{\text{obs},t},\bm{M}\right]$ and applies the loss function to several randomly selected points across the entire set of coordinates. However, three main differences distinguish our 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 foundation: DPF relies on heuristic designs without convergence guarantees, while our approach provides rigorous theoretical analysis. As shown in Appendix [D](https://arxiv.org/html/2509.20098v1#A4), 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 approach—incorporating masks during training to predict clean data—several critical distinctions emerge upon closer examination. First, ambient diffusion lacks a theoretical analysis of how different mask distributions affect the learning dynamics, whereas our work provides a rigorous characterization through Theorem [1](https://arxiv.org/html/2509.20098v1#Thmtheorem1) parts (ii) and (iii). The most significant difference lies in the sampling methodology. Ambient diffusion employs a fixed mask sampling strategy and directly approximates $\mathbb{E}\left[\bm{x}_{0}\mid\bm{M}_{\text{ctx}}\odot\bm{x}_{\text{obs},t},\bm{M}_{\text{ctx}}\right]$ to approximate $\mathbb{E}\left[\bm{x}_{0}\mid\bm{x}_{\text{obs},t},\bm{M}\right]$. As we demonstrate in Theorem [2](https://arxiv.org/html/2509.20098v1#Thmtheorem2), this approximation introduces a 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 to approximate the true conditional expectation, providing theoretical convergence guarantees and eliminating the distribution gap inherent in the ambient diffusion approach.

### C.3 Score matching with missing data

Our approach differs from this prior work (Givens et al., 2025) in a 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 $\mathbb{E}\left[\bm{x}_{0}\mid\bm{x}_{\text{obs},t},\bm{M}\right]$, which requires only the observed data $\bm{x}_{\text{obs},t}$ and mask $\bm{M}$ as inputs. By contrast, noise matching and velocity matching require computing the score term $\nabla_{\bm{x}_{\text{unobs},t}}\log p_{t}\left(\bm{x}_{\text{unobs},t}\mid\bm{x}_{\text{obs}},\bm{M}\right)$, which depends on the unobserved data $\bm{x}_{\text{unobs},t}$ that is unavailable during training. To address this limitation, the score matching approach relies on a 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, we focus exclusively on the data matching framework.

### C.4 Extension to other score-based models

Our approach can be extended to other score-based generative models through established theoretical 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), provides a foundation for this extension. Through Tweedie’s formula (Efron, 2011), we can establish the connection between the score function $\nabla_{\bm{x}}\log p_{t}\left(\bm{x}_{t}\right)$ and the conditional expectation $\mathbb{E}\left[\bm{x}_{0}\mid\bm{x}_{t}\right]$. Leveraging this connection via the parameterization trick, we can train a model to learn $\mathbb{E}\left[\bm{x}_{0}\mid\bm{x}_{\text{obs},t},\bm{M}\right]$ using our proposed method, which can then be converted to output the score function for the denoising process (Zhou et al., 2024).

### C.5 Method comparison and theoretical analysis

We provide additional method comparisons in Tab. [3](https://arxiv.org/html/2509.20098v1#A3.T3) and summarize the key feature comparisons with the most related methods, contrasting our approach with three existing methods: (1) AmbientDiff (Daras et al., 2023), (2) Diffusion Probabilistic Fields (DPF) (Zhuang et al., 2023), and (3) Score Matching with Missing Data (Givens et al., 2025).

**Table 3: Comparison of the most related methods.**
| Aspect | Ours | Ambient Diffusion | DPF | MissDiff |
| --- | --- | --- | --- | --- |
| Training Objective | Data matching | Data matching | Noise matching | Noise matching |
| Model Input | Context only | Fixed mask sampling | Both context and query | Masked tabular data |
| Query Mask Usage | Hidden during training | Hidden during training | Provided to model | Provided to model |
| Mask Sampling | Random subsets: $\bm{M}_{\text{ctx}}\subseteq\bm{M}$ | Fixed distribution | Both $\bm{M}_{\text{ctx}}$ and $\bm{M}_{\text{qry}}$ | $\bm{M}_{\text{ctx}}=\bm{M}_{\text{qry}}=\bm{M}$ |
| Expectation Approx. | Ensemble: $\frac{1}{n}\sum_{i}\bm{x}_{\theta}(t,\text{ctx})$ | Direct: $\mathbb{E}[\bm{x}_{0}\mid\text{ctx}]$ | Incorrect | Direct: $\mathbb{E}[\bm{x}_{0}\mid\text{obs}]$ |
| Theoretical Guarantees | ✓ Convergence proofs (Thm. [1](https://arxiv.org/html/2509.20098v1#Thmtheorem1)& [2](https://arxiv.org/html/2509.20098v1#Thmtheorem2)) | × Lacks rigorous analysis | × Heuristic design | × Lacks rigorous analysis |
| Distribution Gap | ✓ Minimized via ensemble | × Gap $\propto$ Var[conditioning] | × Not addressed | × Not addressed |
| Learning Dynamics | ✓ Gradient scaling analysis | × No analysis | × No analysis | × No analysis |

## Appendix D Method comparison: diffusion probabilistic fields

We summarize the training and sampling algorithms for diffusion probabilistic fields (DPF) (Zhuang et al., 2023) in Alg. [2](https://arxiv.org/html/2509.20098v1#alg2), [3](https://arxiv.org/html/2509.20098v1#alg3). We cannot directly adopt DPF for our 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.

Figure: Algorithm 2 DPF training process (Zhuang et al., 2023)

Figure: Algorithm 3 DPF sampling process (Zhuang et al., 2023)

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 to predict conditional expectation has a more flexible requirement on model input (see Sec. [C.5](https://arxiv.org/html/2509.20098v1#A3.SS5) for details).
- 2.
The DPF method suggests we should take $\bm{M}_{\text{qry}}$ to be the full observed mask, which is impossible to implement in our task.
- 3.
DPF trains a model that takes both ctx and qry as inputs. In contrast, our method trains a model using only ctx as input. This difference leads the model to converge to a 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. [2](https://arxiv.org/html/2509.20098v1#alg2) and demonstrate that it does not yield the desired solution for the denoising process.

We reparameterize the loss function as

$$ $\mathcal{L}\left(\bm{\theta}\right)=\|\bm{M}_{\text{qry}}\odot\left(\bm{x}_{\bm{\theta}}(t,\text{ctx},\text{qry})-\bm{x}_{0}\right)\|^{2}$ (18) $$

When optimized,

$$ $\bm{M}_{\text{qry}}\odot\left(\alpha_{t}\bm{x}_{\bm{\theta}}^{*}\left(t,\text{ctx},\text{qry}\right)+\sigma_{t}\bm{\epsilon}_{\bm{\theta}}^{*}\left(t,\text{ctx},\text{qry}\right)-\text{qry}[1]\right)=\bm{0}$ (19) $$

In the following, we will analyze the optimal solution given by equation [18](https://arxiv.org/html/2509.20098v1#A4.E18). The conditional expectation $\mathbb{E}\left[\bm{x}_{0}\mid\text{ctx},\text{qry}\right]$ minimizes the expected squared error (Bishop & Nasrabadi, 2006). Hence, we have

$$ $\bm{M}_{\text{qry}}\odot\bm{x}_{\bm{\theta}}^{*}(t,\text{ctx},\text{qry})=\bm{M}_{\text{qry}}\odot\mathbb{E}\left[\bm{x}_{0}\mid\text{ctx},\text{qry}\right]$ (20) $$

For simplicity, we consider a simple distribution $\bm{x}_{0}\sim p_{\text{data}}=\mathcal{N}(\bm{\mu},\bm{\Sigma})$. The noisy observations are defined as: $\bm{z}_{\text{ctx}}=\alpha_{t}\bm{x}_{0}+\sigma_{t}\bm{\epsilon}_{\text{ctx}},\bm{z}_{\text{qry}}=\alpha_{t}\bm{x}_{0}+\sigma_{t}\bm{\epsilon}_{\text{qry}}$. We have a joint Gaussian distribution:

$$ $\begin{bmatrix}\bm{x}_{0}\\ \bm{z}_{\text{ctx}}\\ \bm{z}_{\text{qry}}\end{bmatrix}\sim\mathcal{N}\left(\begin{bmatrix}\bm{\mu}\\ \alpha_{t}\bm{\mu}\\ \alpha_{t}\bm{\mu}\end{bmatrix},\begin{bmatrix}\bm{\Sigma}&\alpha_{t}\bm{\Sigma}&\alpha_{t}\bm{\Sigma}\\ \alpha_{t}\bm{\Sigma}&\alpha_{t}^{2}\bm{\Sigma}+\sigma_{t}^{2}\bm{I}&\alpha_{t}^{2}\bm{\Sigma}\\ \alpha_{t}\bm{\Sigma}&\alpha_{t}^{2}\bm{\Sigma}&\alpha_{t}^{2}\bm{\Sigma}+\sigma_{t}^{2}\bm{I}\end{bmatrix}\right)$ (21) $$

Let $\bm{y}$ denote the observed entries selected by the masks: $\bm{y}=\bm{S}\begin{bmatrix}\bm{z}_{\text{ctx}}\\
\bm{z}_{\text{qry}}\end{bmatrix}$, where $\bm{S}$ is the selection matrix that extracts the masked entries from $\bm{z}_{\text{ctx}}$ and $\bm{z}_{\text{qry}}$.
By the property of linear transformation of multivariate Gaussian, suppose we have a joint Gaussian distribution of the form:

$$ $\begin{bmatrix}\bm{x}_{0}\\ \bm{z}\end{bmatrix}\sim\mathcal{N}\left(\begin{bmatrix}\bm{\mu}_{x}\\ \bm{\mu}_{z}\end{bmatrix},\begin{bmatrix}\bm{\Sigma}_{xx}&\bm{\Sigma}_{xz}\\ \bm{\Sigma}_{zx}&\bm{\Sigma}_{zz}\end{bmatrix}\right)$ (22) $$

then the joint distribution of $\bm{x}_{0}$ and $\bm{y}$ is given by

$$ $\begin{bmatrix}\bm{x}_{0}\\ \bm{y}\end{bmatrix}=\begin{bmatrix}\bm{x}_{0}\\ \bm{S}\bm{z}\end{bmatrix}\sim\mathcal{N}\left(\begin{bmatrix}\bm{\mu}_{x}\\ \bm{S}\bm{\mu}_{z}\end{bmatrix},\begin{bmatrix}\bm{\Sigma}_{xx}&\bm{\Sigma}_{xz}\bm{S}^{\top}\\ \bm{S}\bm{\Sigma}_{zx}&\bm{S}\bm{\Sigma}_{zz}\bm{S}^{\top}\end{bmatrix}\right)$ (23) $$

Using the above property, the joint distribution of $\bm{x}_{0}$ and $\bm{y}$ is then:

$$ $\displaystyle\begin{bmatrix}\bm{x}_{0}\\ \bm{y}\end{bmatrix}=\begin{bmatrix}\bm{x}_{0}\\ \bm{S}\begin{bmatrix}\bm{z}_{\text{ctx}}\\ \bm{z}_{\text{qry}}\end{bmatrix}\end{bmatrix}$ (24a) $\displaystyle\sim\mathcal{N}\left(\begin{bmatrix}\bm{\mu}\\ \alpha_{t}\bm{S}\left(\bm{1}_{2}\otimes\bm{I}\right)\bm{\mu}\end{bmatrix},\begin{bmatrix}\bm{\Sigma}_{xx}&\alpha_{t}\left(\bm{1}_{2}^{\top}\otimes\bm{\Sigma}\right)\bm{S}^{\top}\\ \alpha_{t}\bm{S}\left(\bm{1}_{2}\otimes\bm{\Sigma}\right)&\bm{S}\left(\alpha_{t}^{2}\left(\bm{1}_{2}\bm{1}_{2}^{\top}\otimes\bm{\Sigma}\right)+\sigma_{t}^{2}\bm{I}\right)\bm{S}^{\top}\end{bmatrix}\right)$ (24b) $$

Thus, the conditional expectation is:

$$ $\mathbb{E}\left[\bm{x}_{0}\mid\bm{y}\right]=\bm{\mu}+\alpha_{t}\left(\bm{1}_{2}^{\top}\otimes\bm{\Sigma}\right)\bm{S}^{\top}\left[\bm{S}\left(\alpha_{t}^{2}\left(\bm{1}_{2}\bm{1}_{2}^{\top}\otimes\bm{\Sigma}\right)+\sigma_{t}^{2}\bm{I}\right)\bm{S}^{\top}\right]^{-1}\left(\bm{y}-\alpha_{t}\bm{S}\left(\bm{1}_{2}\otimes\bm{I}\right)\bm{\mu}\right)$ (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 a theoretical guarantee: with sufficient sampling of context masks, the model’s output will converge to the desired objective (see Theorem [2](https://arxiv.org/html/2509.20098v1#Thmtheorem2)).

## Appendix E Multi-step sampling

In the following, we will consider two specific settings of the mask $\bm{M}$ and diffusion time $t$, and use these constructions to approximate $\mathbb{E}\left[\bm{x}_{0}\mid\bm{x}_{t},\bm{x}_{\text{obs}},\bm{M}\right]$.

#### Diffusion expectation approximation.

We randomly generate multiple masks $\{\bm{M}_{\text{rnd}}^{(i)}\}_{i=1}^{K}$, where $\bm{M}_{\text{rnd}}$ follows the same marginal distribution as $\bm{M}_{\text{ctx}}$ in the training process, but not necessarily being a subset of $\bm{M}$ and take the average across all samples yields:

$$ $\mathbb{E}\left[\bm{x}_{0}\mid\bm{x}_{t}\right]\approx\frac{1}{K}\sum_{k=1}^{K}\bm{x}_{\bm{\theta}}\left(t,\bm{M}_{\text{rnd}}^{(k)}\odot\bm{x}_{t},\bm{M}_{\text{rnd}}^{(k)}\right)$ (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 us to obtain the same distributional modeling capability as standard diffusion models trained on complete datasets.

#### Imputation expectation approximation.

Given partially observed samples $\bm{x}_{\text{obs}}$, we apply the forward diffusion process to a small timestamp $t=\delta$ and generate random masks $\{\bm{M}_{\text{ctx}}^{(k)}\}_{k=1}^{K}\subseteq\bm{M}$ to approximate:

$$ $\mathbb{E}\left[\bm{x}_{0}\mid\bm{x}_{\text{obs}},\bm{M}\right]\approx\mathbb{E}\left[\bm{x}_{0}\mid\bm{x}_{\text{obs},\delta},\bm{M}\right]\approx\frac{1}{K}\sum_{k=1}^{K}\bm{x}_{\bm{\theta}}\left(\delta,\bm{M}_{\text{ctx}}^{(k)}\odot\bm{x}_{\text{obs},\delta},\bm{M}_{\text{ctx}}^{(k)}\right)$ (27) $$

The optimal denoiser $\mathbb{E}\left[\bm{x}_{0}\mid\bm{x}_{t},\bm{x}_{\text{obs}},\bm{M}\right]$ requires the expectation of $\bm{x}_{0}$ conditional on all three information sources: the noisy state $\bm{x}_{t}$, the clean observations $\bm{x}_{\text{obs}}$, and the observation mask $\bm{M}$. Our approach decomposes this complex conditioning into two manageable components: the *diffusion expectation* $\mathbb{E}\left[\bm{x}_{0}\mid\bm{x}_{t}\right]$ captures the denoising information from the current noisy state, while the *imputation expectation* $\mathbb{E}\left[\bm{x}_{0}\mid\bm{x}_{\text{obs}},\bm{M}\right]$ incorporates the structural information from the observed values and their locations. We then heuristically combine these two sources of information through a weighted average:

$$ $\displaystyle\hat{\bm{x}_{\bm{\theta}}}\left(t,\bm{x}_{t},\bm{x}_{\text{obs}},\bm{M}\right)=\mathbb{E}\left[\bm{x}_{0}\mid\bm{x}_{t},\bm{x}_{\text{obs}},\bm{M}\right]\approx\omega_{t}\mathbb{E}\left[\bm{x}_{0}\mid\bm{x}_{t}\right]+\left(1-\omega_{t}\right)\mathbb{E}\left[\bm{x}_{0}\mid\bm{x}_{\text{obs}},\bm{M}\right]$ (28) $\displaystyle\approx\omega_{t}\mathbb{E}_{\bm{M}_{\text{rnd}}}\left[\bm{x}_{\bm{\theta}}\left(t,\bm{M}_{\text{rnd}}\odot\bm{x}_{t},\bm{M}_{\text{rnd}}\right)\right]+\left(1-\omega_{t}\right)\mathbb{E}_{\bm{M}_{\text{ctx}}\subseteq\bm{M}}\left[\bm{x}_{\bm{\theta}}\left(\delta,\bm{M}_{\text{ctx}}\odot\bm{x}_{\text{obs},\delta},\bm{M}_{\text{ctx}}\right)\right]$ $$

where $\omega_{t}$ is a monotonically increasing weight function that transitions from 0 to $1$ as the diffusion process progresses and $\delta$ is a sufficiently small positive number.

Figure: Algorithm 4 Diffusion-based sampling for data imputation

We present our proposed sampling algorithm in Alg. [4](https://arxiv.org/html/2509.20098v1#alg4). At the implementation level, we precompute the *imputation expectation* once before the denoising process begins. During each denoising step, we approximate the *diffusion expectation* by sampling a single random mask $\bm{M}_{\text{rnd}}$ and making one model evaluation. This single-sample approximation is analogous to using a batch size of 1 in stochastic gradient descent, where we accept the variance from using only one sample in exchange for computational efficiency. Following Lugmayr et al. (2022), our sampling procedure applies different denoising strategies to observed and unobserved regions. For unobserved elements, we estimate the noise using our trained model, while for observed elements, we directly compute the noise from the known clean observations. This approach ensures that the observed values remain consistent with their true underlying data throughout the denoising process. A more advanced setting proposed in Huang et al. (2024) uses guided diffusion sampling that starts from Gaussian noise and iteratively denoises it while being guided by two loss terms: an observation loss (matching sparse measurements) and a PDE loss (satisfying the governing equation), ultimately generating complete solutions that are consistent with both the partial 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.

## Appendix 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 $\bm{x}_{\bm{\theta}}\left(t,\mathrm{ctx}\right)=\mathbb{E}\left[\bm{x}_{0}\mid\mathrm{ctx}\right]+\bm{b}\left(\mathrm{ctx}\right)+\bm{\epsilon}_{\mathrm{bias}}\left(\mathrm{ctx}\right)$. Given a context $\mathrm{ctx}$, the following three components are mutually uncorrelated:

- •
data component: $\mathbb{E}[\bm{x}_{0}\mid\mathrm{ctx}]-\mathbb{E}[\mathbb{E}[\bm{x}_{0}\mid\mathrm{ctx}]]$,
- •
bias component: $\bm{b}(\mathrm{ctx})-\mathbb{E}[\bm{b}(\mathrm{ctx})]$,
- •
random error: $\bm{\epsilon}(\mathrm{ctx})$.

We also assume that $\bm{\epsilon}(\mathrm{ctx}^{(i)})$ and $\bm{\epsilon}(\mathrm{ctx}^{(j)})$ are independent for $i\neq j$.

For notation simplicity, we denote

$$ obs $\displaystyle=\left[\bm{M}\odot\bm{x}_{\text{obs},t},\bm{M}\right]$ (29a) ctx $\displaystyle=\left[\bm{M}_{\text{ctx}}\odot\bm{x}_{\text{obs},t},\bm{M}_{\text{ctx}}\right]$ (29b) qry $\displaystyle=\left[\bm{M}_{\text{qry}}\odot\bm{x}_{\text{obs},t},\bm{M}_{\text{qry}}\right]$ (29c) $$

when it is clear from the context.

###### Assumption 1 (Uncorrelated decomposition) .

### 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 $\bm{X}$, $\bm{Y}$, and $\bm{Z}$ be random vectors in $\mathbb{R}^{d}$ such that the relevant second moments are finite. Let $\bm{g}:\text{Space}(\bm{Y})\to\mathbb{R}^{d}$ be a deterministic function and let the objective function $L(\bm{g})$ be defined as

$$ $L(\bm{g})=\mathbb{E}\left[\|\bm{Z}\odot\bm{g}(\bm{Y})-\bm{Z}\odot\bm{X}\|^{2}\right],$ (30) $$

where $\odot$ denotes the Hadamard (element-wise) product.
If each component of the vector $\mathbb{E}[\bm{Z}\odot\bm{Z}\mid\bm{Y}]$ is strictly positive almost surely, then the unique function $\bm{g}^{*}$ that minimizes $L(\bm{g})$ is given by

$$ $\bm{g}^{*}(\bm{Y})=\frac{\mathbb{E}[\bm{Z}\odot\bm{Z}\odot\bm{X}\mid\bm{Y}]}{\mathbb{E}[\bm{Z}\odot\bm{Z}\mid\bm{Y}]},$ (31) $$

where the division is performed element-wise.

###### Proof.

The objective function $L(\bm{g})$ can be decomposed by writing the squared Euclidean norm as a sum over its components.

$$ $\displaystyle L(\bm{g})$ $\displaystyle=\mathbb{E}\left[\|\bm{Z}\odot(\bm{g}(\bm{Y})-\bm{X})\|^{2}\right]$ (32a) $\displaystyle=\mathbb{E}\left[\sum_{i=1}^{d}\left(\bm{Z}_{i}(\bm{g}_{i}(\bm{Y})-\bm{X}_{i})\right)^{2}\right]$ (32b) $\displaystyle=\sum_{i=1}^{d}\mathbb{E}\left[\bm{Z}_{i}^{2}(\bm{g}_{i}(\bm{Y})-\bm{X}_{i})^{2}\right]$ (32c) $$

The final step follows from the linearity of expectation. The total loss is a sum of non-negative terms, so $L(\bm{g})$ is minimized if and only if each term in the summation is minimized independently. Let $L_{i}(\bm{g}_{i})=\mathbb{E}[\bm{Z}_{i}^{2}(\bm{g}_{i}(\bm{Y})-\bm{X}_{i})^{2}]$ be the $i$-th term. The optimization problem is thus reduced to finding the function $\bm{g}_{i}$ that minimizes $L_{i}$ for each component $i\in\{1,\dots,d\}$.
By the law of total expectation, $L_{i}(\bm{g}_{i})$ can be written as:

$$ $L_{i}(\bm{g}_{i})=\mathbb{E}_{\bm{Y}}\left[\mathbb{E}\left[\bm{Z}_{i}^{2}(\bm{g}_{i}(\bm{Y})-\bm{X}_{i})^{2}\mid\bm{Y}\right]\right]$ (33) $$

The outer expectation is minimized by minimizing the inner conditional expectation for any given realization $\bm{y}$ from the space of $\bm{Y}$. For a fixed $\bm{y}$, let $\bm{v}_{i}=\bm{g}_{i}(\bm{y})$ be a deterministic scalar. The inner expectation becomes:

$$ $\mathbb{E}\left[\bm{Z}_{i}^{2}(\bm{v}_{i}-\bm{X}_{i})^{2}\mid\bm{Y}=\bm{y}\right]=\mathbb{E}\left[\bm{Z}_{i}^{2}(\bm{v}_{i}^{2}-2\bm{v}_{i}\bm{X}_{i}+\bm{X}_{i}^{2})\mid\bm{Y}=\bm{y}\right]$ (34) $$

Applying the linearity of conditional expectation yields a quadratic function of $\bm{v}_{i}$:

$$ $\bm{v}_{i}^{2}\mathbb{E}[\bm{Z}_{i}^{2}\mid\bm{Y}=\bm{y}]-2\bm{v}_{i}\mathbb{E}[\bm{Z}_{i}^{2}\bm{X}_{i}\mid\bm{Y}=\bm{y}]+\mathbb{E}[\bm{Z}_{i}^{2}\bm{X}_{i}^{2}\mid\bm{Y}=\bm{y}]$ (35) $$

This is a convex quadratic in $\bm{v}_{i}$, since its leading coefficient $\mathbb{E}[\bm{Z}_{i}^{2}\mid\bm{Y}=\bm{y}]$ is strictly positive by hypothesis. The unique minimum is found by setting the derivative with respect to $\bm{v}_{i}$ to zero:

$$ $2\bm{v}_{i}\mathbb{E}[\bm{Z}_{i}^{2}\mid\bm{Y}=\bm{y}]-2\mathbb{E}[\bm{Z}_{i}^{2}\bm{X}_{i}\mid\bm{Y}=\bm{y}]=0$ (36) $$

Solving for $\bm{v}_{i}$ gives the optimal value for the component function at $\bm{y}$:

$$ $\bm{v}_{i}=\frac{\mathbb{E}[\bm{Z}_{i}^{2}\bm{X}_{i}\mid\bm{Y}=\bm{y}]}{\mathbb{E}[\bm{Z}_{i}^{2}\mid\bm{Y}=\bm{y}]}$ (37) $$

This establishes the optimal form for each component $\bm{g}_{i}^{*}(\bm{y})$ of the function $\bm{g}^{*}(\bm{y})$. Since this holds for all $\bm{y}$, the optimal function $\bm{g}_{i}^{*}$ for the $i$-th component is:

$$ $\bm{g}_{i}^{*}(\bm{Y})=\frac{\mathbb{E}[\bm{Z}_{i}^{2}\bm{X}_{i}\mid\bm{Y}]}{\mathbb{E}[\bm{Z}_{i}^{2}\mid\bm{Y}]}$ (38) $$

Assembling the components for $i=1,\dots,d$ into a single vector equation gives the expression for the optimal function $\bm{g}^{*}$:

$$ $\bm{g}^{*}(\bm{Y})=\frac{\mathbb{E}[\bm{Z}\odot\bm{Z}\odot\bm{X}\mid\bm{Y}]}{\mathbb{E}[\bm{Z}\odot\bm{Z}\mid\bm{Y}]}$ (39) $$

where the division is understood to be element-wise.
∎

Having established the fundamental relationship between optimal model outputs and query probabilities in the lemma, we now proceed to prove our main theorem.

See [1](https://arxiv.org/html/2509.20098v1#Thmtheorem1)

###### Proof.

Given the training algorithm in Alg. [1](https://arxiv.org/html/2509.20098v1#alg1), we have

$$ $\displaystyle\bm{\theta}^{*}$ $\displaystyle=\operatorname*{arg\,min}_{\bm{\theta}}\mathbb{E}_{t,\bm{x}_{\text{obs}},\left(\bm{M}_{\text{ctx}},\bm{M}_{\text{qry}}\subseteq\bm{M}\right)}\left[\|\bm{M}_{\text{qry}}\odot\left(\bm{x}_{\bm{\theta}}\left(t,\text{ctx}\right)-\bm{x}_{\text{obs}}\right)\|^{2}\right]$ (40a) $\displaystyle=\operatorname*{arg\,min}_{\bm{\theta}}\mathbb{E}_{t}\mathbb{E}_{\bm{x}_{\text{obs}},\left(\bm{M}_{\text{ctx}},\bm{M}_{\text{qry}}\subseteq\bm{M}\right)}\left[\|\bm{M}_{\text{qry}}\odot\left(\bm{x}_{\bm{\theta}}\left(t,\text{ctx}\right)-\bm{M}\odot\bm{x}_{0}\right)\|^{2}\right]$ (40b) $\displaystyle=\operatorname*{arg\,min}_{\bm{\theta}}\mathbb{E}_{t}\mathbb{E}_{\bm{x}_{\text{obs}},\left(\bm{M}_{\text{ctx}},\bm{M}_{\text{qry}}\subseteq\bm{M}\right)}\left[\|\bm{M}_{\text{qry}}\odot\bm{x}_{\bm{\theta}}\left(t,\text{ctx}\right)-\bm{M}_{\text{qry}}\odot\bm{x}_{0}\|^{2}\right]$ (40c) $$

(i) Optimal solution: When optimized, given $\forall\bm{M}_{\text{ctx}}\subseteq\bm{M}$, for any index $i$ such that $P((\bm{M}_{\text{qry}})_{i}=1\mid\bm{M}_{\text{ctx}})>0$ under the sampling distribution, we have

$$ $\bm{M}_{\text{qry}}\odot\bm{x}_{\bm{\theta}}\left(t,\text{ctx}\right)=\mathbb{E}\left[\bm{M}_{\text{qry}}\odot\bm{x}_{0}\mid\text{ctx},\bm{M}_{\text{qry}}\right]=\bm{M}_{\text{qry}}\odot\mathbb{E}\left[\bm{x}_{0}\mid\text{ctx}\right]$ (41) $$

Thus, $\bm{M}_{\text{qry}}\odot\left(\bm{x}_{\bm{\theta}}\left(t,\text{ctx}\right)-\mathbb{E}\left[\bm{x}_{0}\mid\text{ctx}\right]\right)=\bm{0}$ for any $\bm{M}_{\text{qry}}$ in the support of the sampling distribution given $\bm{M}_{\text{ctx}}$.

Case 1: When $P((\bm{M}_{\text{qry}})_{i}=1\mid\bm{M}_{\text{ctx}})>0$. By applying Lemma [1](https://arxiv.org/html/2509.20098v1#Thmlemma1) with $\bm{Z}=\bm{M}_{\text{qry}}$, $\bm{Y}=\text{ctx}$, $\bm{X}=\bm{x}_{0}$, and $\bm{g}(\cdot)=\bm{x}_{\bm{\theta}}(t,\cdot)$, the optimal solution is:

$$ $\bm{x}_{\bm{\theta}}^{*}(t,\text{ctx})=\frac{\mathbb{E}[\bm{M}_{\text{qry}}\odot\bm{M}_{\text{qry}}\odot\bm{x}_{0}\mid\text{ctx}]}{\mathbb{E}[\bm{M}_{\text{qry}}\odot\bm{M}_{\text{qry}}\mid\text{ctx}]}$ (42) $$

Since $\bm{M}_{\text{qry}}$ is a binary mask where $(\bm{M}_{\text{qry}})_{i}\in\{0,1\}$, we have $(\bm{M}_{\text{qry}})_{i}\odot(\bm{M}_{\text{qry}})_{i}=(\bm{M}_{\text{qry}})_{i}$. Thus:

$$ $\bm{x}_{\bm{\theta}}^{*}(t,\text{ctx})=\frac{\mathbb{E}[\bm{M}_{\text{qry}}\odot\bm{x}_{0}\mid\text{ctx}]}{\mathbb{E}[\bm{M}_{\text{qry}}\mid\text{ctx}]}$ (43) $$

For component $i$ where $P((\bm{M}_{\text{qry}})_{i}=1\mid\bm{M}_{\text{ctx}})>0$:

$$ $\displaystyle(\bm{x}_{\bm{\theta}}^{*}(t,\text{ctx}))_{i}$ $\displaystyle=\frac{\mathbb{E}[(\bm{M}_{\text{qry}})_{i}(\bm{x}_{0})_{i}\mid\text{ctx}]}{\mathbb{E}[(\bm{M}_{\text{qry}})_{i}\mid\text{ctx}]}$ (44a) $\displaystyle=\frac{\mathbb{E}[(\bm{M}_{\text{qry}})_{i}(\bm{x}_{0})_{i}\mid\text{ctx}]}{P((\bm{M}_{\text{qry}})_{i}=1\mid\bm{M}_{\text{ctx}})}$ (44b) $\displaystyle=\frac{\mathbb{E}[(\bm{x}_{0})_{i}\mathbf{1}_{(\bm{M}_{\text{qry}})_{i}=1}\mid\text{ctx}]}{P((\bm{M}_{\text{qry}})_{i}=1\mid\bm{M}_{\text{ctx}})}$ (44c) $\displaystyle=\mathbb{E}[(\bm{x}_{0})_{i}\mid\text{ctx},(\bm{M}_{\text{qry}})_{i}=1]$ (44d) $\displaystyle=\mathbb{E}[(\bm{x}_{0})_{i}\mid\text{ctx}]$ (44e) $$

where the last equality holds because $\bm{x}_{0}$ is independent of $\bm{M}_{\text{qry}}$ given ctx.

Case 2: When $P((\bm{M}_{\text{qry}})_{i}=1\mid\bm{M}_{\text{ctx}})=0$. In this case, $(\bm{M}_{\text{qry}})_{i}=0$ almost surely given $\bm{M}_{\text{ctx}}$. The contribution of index $i$ to the loss function is:

$$ $\displaystyle\quad\mathbb{E}[|(\bm{M}_{\text{qry}})_{i}(\bm{x}_{\bm{\theta}}(t,\text{ctx}))_{i}-(\bm{M}_{\text{qry}})_{i}(\bm{x}_{0})_{i}|^{2}\mid\text{ctx}]$ $\displaystyle=\mathbb{E}[|0\cdot(\bm{x}_{\bm{\theta}}(t,\text{ctx}))_{i}-0\cdot(\bm{x}_{0})_{i}|^{2}\mid\text{ctx}]$ (45a) $\displaystyle=0$ (45b) $$

Therefore, $(\bm{x}_{\bm{\theta}}(t,\text{ctx}))_{i}$ does not affect the loss function and can take any arbitrary value.

Thus, we have

$$ $\left(\bm{x}_{\bm{\theta}}\left(t,\text{ctx}\right)\right)_{i}=\begin{cases}\mathbb{E}\left[\left(\bm{x}_{0}\right)_{i}\mid\text{ctx}\right],&P(\left(\bm{M}_{\text{qry}}\right)_{i}=1\mid\bm{M}_{\text{ctx}})>0\\ \text{an arbitrary value},&P(\left(\bm{M}_{\text{qry}}\right)_{i}=1\mid\bm{M}_{\text{ctx}})=0\end{cases}$ (46) $$

If the union of all possible query mask $\bm{M}_{\text{qry}}$ supports covers all feature dimensions, we have $\bigcup_{\bm{M}_{\text{qry}}\text{ possible}}\text{supp}(\bm{M}_{\text{qry}})=\{1,\ldots,\dim(\bm{x}_{0})\}$. Thus, $\forall i,P(\left(\bm{M}_{\text{qry}}\right)_{i}=1)>0$ and

$$ $\bm{x}_{\bm{\theta}}\left(t,\text{ctx}\right)=\mathbb{E}\left[\bm{x}_{0}\mid\text{ctx}\right]$ (47) $$

(ii) Gradient magnitude scaling: The gradient of the loss with respect to the $i$-th output component is:

$$ $\frac{\partial\mathcal{L}}{\partial(\bm{x}_{\bm{\theta}})_{i}}=2(\bm{M}_{\text{qry}})_{i}\cdot((\bm{x}_{\bm{\theta}})_{i}-(\bm{x}_{\text{obs}})_{i})$ (48) $$

Taking expectation over the sampling distribution:

$$ $\displaystyle\mathbb{E}\left[\left(\frac{\partial\mathcal{L}}{\partial(\bm{x}_{\bm{\theta}})_{i}}\right)^{2}\right]$ $\displaystyle=\mathbb{E}\left[4(\bm{M}_{\text{qry}})_{i}^{2}\cdot((\bm{x}_{\bm{\theta}})_{i}-(\bm{x}_{\text{obs}})_{i})^{2}\right]$ (49a) $\displaystyle=4\mathbb{E}\left[(\bm{M}_{\text{qry}})_{i}\cdot((\bm{x}_{\bm{\theta}})_{i}-(\bm{x}_{\text{obs}})_{i})^{2}\right]$ (49b) $\displaystyle=4p_{i}\mathbb{E}\left[((\bm{x}_{\bm{\theta}})_{i}-(\bm{x}_{\text{obs}})_{i})^{2}\Big|(\bm{M}_{\text{qry}})_{i}=1\right]$ (49c) $$

Let $C_{i}:=\mathbb{E}\left[((\bm{x}_{\bm{\theta}})_{i}-(\bm{x}_{\text{obs}})_{i})^{2}\Big|(\bm{M}_{\text{qry}})_{i}=1\right]$. Then:

$$ $\mathbb{E}\left[\left(\frac{\partial\mathcal{L}}{\partial(\bm{x}_{\bm{\theta}})_{i}}\right)^{2}\right]=4p_{i}C_{i}$ (50) $$

establishing linear scaling with $p_{i}$.

(iii) Parameter update frequency: At each training step, the gradient contribution from feature $i$ is:

$$ $\frac{\partial\mathcal{L}}{\partial\bm{\theta}}\Big|_{\text{feature }i}=\frac{\partial\mathcal{L}}{\partial(\bm{x}_{\bm{\theta}})_{i}}\frac{\partial(\bm{x}_{\bm{\theta}})_{i}}{\partial\bm{\theta}}$ (51) $$

This contribution is non-zero if and only if $(\bm{M}_{\text{qry}})_{i}=1$, which occurs with probability $p_{i}$. Therefore:

$$ $P(\text{feature }i\text{ contributes to parameter update})=P((\bm{M}_{\text{qry}})_{i}=1\mid\bm{M}_{\text{ctx}})=p_{i}$ (52) $$

∎

###### Lemma 1 (Optimal Function for Element-wise Weighted MSE) .

###### Proof.

###### Proof.

### F.2 Partially observed mask convergence theorem

See [2](https://arxiv.org/html/2509.20098v1#Thmtheorem2)

###### Proof.

Define random variables $\bm{\mu}=\mathbb{E}\left[\bm{x}_{0}\mid\text{obs}\right]$ and $\bm{\mu}_{\text{c}}=\mathbb{E}\left[\bm{x}_{0}\mid\text{ctx}\right]$. Our goal is to compute $\mathbb{E}\left[\|\bm{\mu}-\bm{\mu}_{\text{c}}\|^{2}\right]$.

Applying the Law of Total Variance, we have

$$ $\displaystyle\mathrm{Var}\left[\bm{x}_{0}\mid\text{ctx}\right]$ $\displaystyle=\mathbb{E}\left[\mathrm{Var}\left[\bm{x}_{0}\mid\text{obs}\right]\mid\text{ctx}\right]+\mathrm{Var}\left[\mathbb{E}\left[\bm{x}_{0}\mid\text{obs}\right]\mid\text{ctx}\right]$ (53a) $\displaystyle=\mathbb{E}\left[\mathrm{Var}\left[\bm{x}_{0}\mid\text{obs}\right]\mid\text{ctx}\right]+\mathrm{Var}\left[\bm{\mu}\mid\text{ctx}\right]$ (53b) $$

Rearranging:

$$ $\mathrm{Var}\left[\bm{\mu}\mid\text{ctx}\right]=\mathrm{Var}\left[\bm{x}_{0}\mid\text{ctx}\right]-\mathbb{E}\left[\mathrm{Var}\left[\bm{x}_{0}\mid\text{obs}\right]\mid\text{ctx}\right]$ (54) $$

Then:

$$ $\displaystyle\qquad\mathbb{E}\left[\|\bm{\mu}-\bm{\mu}_{\text{c}}\|^{2}\right]$ $\displaystyle=\mathbb{E}\left[\mathbb{E}\left[\|\bm{\mu}-\bm{\mu}_{\text{c}}\|^{2}\mid\text{ctx}\right]\right]$ Law of total expectation (55a) $\displaystyle=\mathbb{E}\left[\mathbb{E}\left[\|\bm{\mu}-\mathbb{E}\left[\bm{x}_{0}\mid\text{ctx}\right]\|^{2}\mid\text{ctx}\right]\right]$ $\displaystyle\text{By definition of }\bm{\mu}_{\text{c}}$ (55b) $\displaystyle=\mathbb{E}\left[\mathbb{E}\left[\|\bm{\mu}-\mathbb{E}\left[\mathbb{E}\left[\bm{x}_{0}\mid\text{obs}\right]\mid\text{ctx}\right]\|^{2}\mid\text{ctx}\right]\right]$ Tower rule (55c) $\displaystyle=\mathbb{E}\left[\mathbb{E}\left[\|\bm{\mu}-\mathbb{E}\left[\bm{\mu}\mid\text{ctx}\right]\|^{2}\mid\text{ctx}\right]\right]$ $\displaystyle\text{By definition of }\bm{\mu}$ (55d) $\displaystyle=\mathbb{E}\left[\mathrm{Var}\left[\bm{\mu}\mid\text{ctx}\right]\right]$ Definition of conditional variance (55e) $\displaystyle=\mathbb{E}\left[\mathrm{Var}\left[\bm{x}_{0}\mid\text{ctx}\right]-\mathbb{E}\left[\mathrm{Var}\left[\bm{x}_{0}\mid\text{obs}\right]\mid\text{ctx}\right]\right]$ equation 54 (55f) $\displaystyle=\mathbb{E}\left[\mathrm{Var}\left[\bm{x}_{0}\mid\text{ctx}\right]\right]-\mathbb{E}\left[\mathrm{Var}\left[\bm{x}_{0}\mid\text{obs}\right]\right]$ Linearity of expectation (55g) $$

Further define $\bm{\mu}_{k}=\bm{x}_{\bm{\theta}}\left(t,\text{ctx}^{(k)}\right)=\mathbb{E}\left[\bm{x}_{0}\mid\text{ctx}^{(k)}\right]+\bm{b}\left(\text{ctx}^{(k)}\right)+\bm{\epsilon}_{\text{bias}}\left(\text{ctx}^{(k)}\right)$, where $\bm{b}$ is the systematic bias and $\bm{\epsilon}_{\text{bias}}$ is the random error with $\mathbb{E}\left[\bm{\epsilon}_{\text{bias}}\right]=\bm{0}$. The ensemble average is:

$$ $\hat{\bm{\mu}}_{K}=\frac{1}{K}\sum_{k=1}^{K}\bm{\mu}_{k}=\frac{1}{K}\sum_{k=1}^{K}\mathbb{E}[\bm{x}_{0}\mid\text{ctx}^{(k)}]+\frac{1}{K}\sum_{k=1}^{K}\bm{b}(\text{ctx}^{(k)})+\frac{1}{K}\sum_{k=1}^{K}\bm{\epsilon}(\text{ctx}^{(k)})$ (56) $$

Computing the expectation:

$$ $\displaystyle\mathbb{E}[\hat{\bm{\mu}}_{K}]$ $\displaystyle=\mathbb{E}\left[\frac{1}{K}\sum_{k=1}^{K}\bm{\mu}_{k}\right]$ (57a) $\displaystyle=\frac{1}{K}\sum_{k=1}^{K}\mathbb{E}[\mathbb{E}[\bm{x}_{0}\mid\text{ctx}^{(k)}]+\bm{b}(\text{ctx}^{(k)})+\bm{\epsilon}(\text{ctx}^{(k)})]$ (57b) $\displaystyle=\mathbb{E}[\mathbb{E}[\bm{x}_{0}\mid\text{ctx}]]+\mathbb{E}[\bm{b}(\text{ctx})]+\mathbb{E}[\bm{\epsilon}(\text{ctx})]$ (57c) $\displaystyle=\mathbb{E}[\mathbb{E}[\bm{x}_{0}\mid\text{ctx}]]+\mathbb{E}[\bm{b}(\text{ctx})]$ (57d) $$

The bias of the ensemble estimator is:

$$ $\operatorname{Bias}(\hat{\bm{\mu}}_{K})=\mathbb{E}[\hat{\bm{\mu}}_{K}]-\mathbb{E}[\bm{x}_{0}\mid\text{obs}]=\mathbb{E}[\mathbb{E}[\bm{x}_{0}\mid\text{ctx}]]-\mathbb{E}[\bm{x}_{0}\mid\text{obs}]+\mathbb{E}[\bm{b}(\text{ctx})]$ (58) $$

For the variance, decompose each component:

$$ $\displaystyle\bm{\mu}_{k}-\mathbb{E}[\bm{\mu}_{k}]$ $\displaystyle=\mathbb{E}[\bm{x}_{0}\mid\text{ctx}^{(k)}]+\bm{b}(\text{ctx}^{(k)})+\bm{\epsilon}(\text{ctx}^{(k)})-\mathbb{E}[\mathbb{E}[\bm{x}_{0}\mid\text{ctx}]]-\mathbb{E}[\bm{b}(\text{ctx})]$ (59a) $\displaystyle=\underbrace{(\mathbb{E}[\bm{x}_{0}\mid\text{ctx}^{(k)}]-\mathbb{E}[\mathbb{E}[\bm{x}_{0}\mid\text{ctx}]])}_{\text{data variation}}+\underbrace{(\bm{b}(\text{ctx}^{(k)})-\mathbb{E}[\bm{b}(\text{ctx})])}_{\text{bias variation}}+\underbrace{\bm{\epsilon}(\text{ctx}^{(k)})}_{\text{random error}}$ (59b) $$

The variance of the ensemble average:

$$ $\displaystyle\mathrm{Var}(\hat{\bm{\mu}}_{K})$ $\displaystyle=\mathbb{E}\left[\left\|\hat{\bm{\mu}}_{K}-\mathbb{E}[\hat{\bm{\mu}}_{K}]\right\|^{2}\right]$ (60a) $\displaystyle=\mathbb{E}\left[\left\|\frac{1}{K}\sum_{k=1}^{K}(\bm{\mu}_{k}-\mathbb{E}[\bm{\mu}_{k}])\right\|^{2}\right]$ (60b) $\displaystyle=\frac{1}{K^{2}}\mathbb{E}\left[\sum_{k=1}^{K}\|\bm{\mu}_{k}-\mathbb{E}[\bm{\mu}_{k}]\|^{2}+2\sum_{i<j}\langle\bm{\mu}_{i}-\mathbb{E}[\bm{\mu}_{i}],\bm{\mu}_{j}-\mathbb{E}[\bm{\mu}_{j}]\rangle\right]$ (60c) $$

Since $\text{ctx}^{(1)},\ldots,\text{ctx}^{(K)}$ are conditionally independent given $\bm{x}_{\text{obs},t}$ and $\bm{M}$, the data variations and bias variations are independent across different $k$. Additionally, under the assumption that $\bm{\epsilon}(\text{ctx}^{(i)})$ and $\bm{\epsilon}(\text{ctx}^{(j)})$ are independent for $i\neq j$:

$$ $\mathbb{E}[\langle\bm{\mu}_{i}-\mathbb{E}[\bm{\mu}_{i}],\bm{\mu}_{j}-\mathbb{E}[\bm{\mu}_{j}]\rangle]=0\quad\forall i\neq j$ (61) $$

Therefore:

$$ $\displaystyle\mathrm{Var}(\hat{\bm{\mu}}_{K})$ $\displaystyle=\frac{1}{K^{2}}\sum_{k=1}^{K}\mathbb{E}[\|\bm{\mu}_{k}-\mathbb{E}[\bm{\mu}_{k}]\|^{2}]$ (62a) $\displaystyle=\frac{1}{K}\mathbb{E}\left[\left\|(\mathbb{E}[\bm{x}_{0}\mid\text{ctx}]-\mathbb{E}[\mathbb{E}[\bm{x}_{0}\mid\text{ctx}]])+(\bm{b}(\text{ctx})-\mathbb{E}[\bm{b}(\text{ctx})])+\bm{\epsilon}(\text{ctx})\right\|^{2}\right]$ (62b) $\displaystyle=\frac{1}{K}\left(\mathrm{Var}[\mathbb{E}[\bm{x}_{0}\mid\text{ctx}]]+\mathrm{Var}[\bm{b}(\text{ctx})]+\mathrm{Var}[\bm{\epsilon}]\right)$ (62c) $$

Here, we used the independence between the three components to separate the variances. Combining bias and variance using the bias-variance decomposition:

$$ $\displaystyle\qquad\mathbb{E}\left[\left\|\hat{\bm{\mu}}_{K}-\mathbb{E}[\bm{x}_{0}\mid\text{obs}]\right\|^{2}\right]$ $\displaystyle=\|\operatorname{Bias}(\hat{\bm{\mu}}_{K})\|^{2}+\mathrm{Var}(\hat{\bm{\mu}}_{K})$ (63a) $\displaystyle=\left\|\mathbb{E}[\mathbb{E}[\bm{x}_{0}\mid\text{ctx}]]-\mathbb{E}[\bm{x}_{0}\mid\text{obs}]+\mathbb{E}[\bm{b}(\text{ctx})]\right\|^{2}$ $\displaystyle\qquad+\frac{1}{K}\left(\mathrm{Var}[\mathbb{E}[\bm{x}_{0}\mid\text{ctx}]]+\mathrm{Var}[\bm{b}(\text{ctx})]+\mathrm{Var}[\bm{\epsilon}]\right)$ (63b) $$

Taking the limit as $K\to\infty$:

$$ $\lim_{K\to\infty}\mathbb{E}\left[\left\|\hat{\bm{\mu}}_{K}-\mathbb{E}[\bm{x}_{0}\mid\text{obs}]\right\|^{2}\right]=\left\|\mathbb{E}[\mathbb{E}[\bm{x}_{0}\mid\text{ctx}]]-\mathbb{E}[\bm{x}_{0}\mid\text{obs}]+\mathbb{E}[\bm{b}(\text{ctx})]\right\|^{2}$ (64) $$

This establishes that:

- •
The average bias $\mathbb{E}[\bm{b}(\text{ctx})]$ across all contexts is not reduced by ensemble averaging.
- •
The variance of the context-dependent bias $\mathrm{Var}[\bm{b}(\text{ctx})]$ is reduced by a factor of $1/K$.
- •
Both data variance and random error variance are also reduced by $1/K$.
- •
The asymptotic error includes both the data bias and the model’s average bias.

∎

###### Proof.

## Appendix G Supplementary Experiments

### G.1 Dataset settings

#### Shallow Water and Advection datasets.

We evaluate our approach on two fundamental geophysical PDE systems with distinct characteristics:

$$ Shallow Water: $\displaystyle\partial_{t}u=fv-g\partial_{x}h,\quad\partial_{t}v=-fu-g\partial_{y}h,\quad\partial_{t}h=-H(\partial_{x}u+\partial_{y}v)$ (65a) Advection: $\displaystyle\partial_{t}u(t,x)+\beta\partial_{x}u(t,x)=0$ (65b) $$

where for the shallow water system, $u$ and $v$ represent velocity components, $h$ denotes height field, $f$ is the Coriolis parameter, $g$ is gravitational acceleration, and $H$ is mean depth. For the advection equation, $u$ represents a scalar field being transported and $\beta$ is the advection velocity. We generate 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\times 32$ spatial resolution and 50 temporal frames. We generate synthetic solutions by randomly sampling physical parameters and initial conditions.
For evaluation, we employ complementary metrics to comprehensively assess reconstruction quality across different physical aspects. For the Shallow Water dataset, we evaluate PDE feasibility loss, which measures how well the reconstructed 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 reconstructed 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.

We evaluate our approach on the incompressible Navier-Stokes equations for isotropic turbulence. The governing equations are:

$$ $\displaystyle\partial_{t}\bm{u}+(\bm{u}\cdot\nabla)\bm{u}$ $\displaystyle=-\nabla p+\nu\nabla^{2}\bm{u}+\bm{f},$ (66a) $\displaystyle\nabla\cdot\bm{u}$ $\displaystyle=0$ (66b) $$

where $\bm{u}=(u,v)$ represents the velocity field, $p$ is pressure, $\nu$ is kinematic viscosity, and $\bm{f}$ is external forcing. The data are generated using either pseudo-spectral solvers with 4th-order Runge-Kutta or 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 a spatial resolution of $64\times 64$ and temporal sequences of 100 frames.
For evaluation, the model generates complete field reconstructions. Performance is measured by computing MSE between reconstructed fields and ground truth across all spatial locations and time steps.

#### ERA5 dataset.

We evaluate our approach on the ERA5 reanalysis dataset, which provides comprehensive atmospheric and surface meteorological variables. ERA5 represents the fifth generation atmospheric reanalysis from the European Centre for Medium-Range Weather Forecasts (ECMWF), combining model data with observations through data assimilation to produce a globally complete and consistent dataset. We utilize one year of hourly data sampled at a spatial resolution of $103\times 120$ grid points in latitude-longitude coordinates, creating temporal sequences of dimension $(365\times 24,103,120)$. The data are segmented into 3-hour windows, and we 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.

### G.2 Ablation study

We conduct ablation studies to validate the effectiveness of key components in our proposed method.

#### Test-time gap introduced by replacing 𝑴 ctx \bm{M}_{\text{ctx}} with 𝑴 \bm{M} .

Our sampling procedure requires multiple context masks $\bm{M}_{\text{ctx}}$ to estimate $\mathbb{E}\left[\bm{x}_{0}\mid\bm{x}_{\text{obs},t},\bm{M}\right]\approx\frac{1}{K}\sum_{k=1}^{K}\bm{x}_{\bm{\theta}}\left(t,\bm{M}_{\text{ctx}}^{(k)}\odot\bm{x}_{t},\bm{M}_{\text{ctx}}^{(k)}\right)$. This ablation study compares our method against directly computing $\mathbb{E}\left[\bm{x}_{0}\mid\bm{x}_{\text{obs},t},\bm{M}\right]$ using $\bm{x}_{\bm{\theta}}\left(t,\bm{M}\odot\bm{x}_{t},\bm{M}\right)$. The direct approach creates a distributional mismatch: during training, the model’s input mask follows the distribution of $\bm{M}_{\text{ctx}}$, but during sampling, the input becomes $\bm{M}$. This mismatch degrades model performance. Tab. [4](https://arxiv.org/html/2509.20098v1#A7.T4) presents experimental results comparing both methods.

**Table 4: Performance comparison of two approaches: (1) imputation with multiple time sampling of $\bm{M}_{\text{ctx}}$ followed by ensemble prediction (Theorem [2](https://arxiv.org/html/2509.20098v1#Thmtheorem2)), versus (2) directly using $\bm{M}$ as $\bm{M}_{\text{ctx}}$, which creates a distributional mismatch between training and testing inputs.**
| Method | Shallow Water | Advection |  |  |
| --- | --- | --- | --- | --- |
| 80% | 30% | 80% | 30% |  |
| Ours ($\bm{M}$) | 2.3983 $\pm$ 0.7880 | 2.6717 $\pm$ 1.4731 | 0.1320 $\pm$ 0.0155 | 0.1655 $\pm$ 0.0537 |
| Ours ($\bm{M}_{\text{ctx}}$) | 0.1878 $\pm$ 0.0054 | 0.7379 $\pm$ 0.1101 | 0.1035 $\pm$ 0.0008 | 0.1189 $\pm$ 0.0069 |

#### Backbone architecture.

To demonstrate the generalizability of our method across different architectures, we evaluate both our proposed approach and baseline methods using two distinct backbones: 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 dimension. Results are presented in Tab. [5](https://arxiv.org/html/2509.20098v1#A7.T5). 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 5: Performance comparison across backbone architectures. Results for our method and baselines using Karras UNet (Karras et al., 2024) and FNO (Li et al., 2020) backbones across two PDE datasets.**
| Method | Backbone | Shallow Water | Advection |  |  |
| --- | --- | --- | --- | --- | --- |
| 80% | 30% | 80% | 30% |  |  |
| MissDiff | UNet | 0.3963 $\pm$ 0.0617 | 1.2570 $\pm$ 0.2146 | 0.1030 $\pm$ 0.0004 | 0.1197 $\pm$ 0.0096 |
| FNO | 0.2917 $\pm$ 0.1683 | 0.7525 $\pm$ 0.1529 | 0.1375 $\pm$ 0.0063 | 0.4816 $\pm$ 0.0187 |  |
| Ours | UNet | 0.3279 $\pm$ 0.0655 | 0.9292 $\pm$ 0.1963 | 0.1035 $\pm$ 0.0008 | 0.1189 $\pm$ 0.0069 |
| FNO | 0.1869 $\pm$ 0.0015 | 0.7379 $\pm$ 0.1101 | 0.1240 $\pm$ 0.0040 | 0.3527 $\pm$ 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 are presented in Table [6](https://arxiv.org/html/2509.20098v1#A7.T6). 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 a direct comparison point that validates our experimental setup.

**Table 6: Performance comparison of context and query mask ratio.**
| Context Ratio | Query Ratio | Navier-Stokes |  |  |
| --- | --- | --- | --- | --- |
| 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 |

#### Optimal denoiser approximation.

We approximate the optimal denoiser $\mathbb{E}\left[\bm{x}_{0}\mid\bm{x}_{t},\bm{x}_{\text{obs}},\bm{M}\right]$ through a weighted combination of diffusion expectation $\mathbb{E}\left[\bm{x}_{0}\mid\bm{x}_{t}\right]$ and imputation expectation $\mathbb{E}\left[\bm{x}_{0}\mid\bm{x}_{\text{obs}},\bm{M}\right]$ using empirical weight $\omega_{t}$ (equation [28](https://arxiv.org/html/2509.20098v1#A5.E28)). We investigate different weighting strategies to understand their impact on reconstruction quality during the multi-step generation process (200 steps). The results can be seen in Tab. [7](https://arxiv.org/html/2509.20098v1#A7.T7).

**Table 7: Impact of weighting strategies on optimal denoiser approximation.**
| Method | Navier-Stokes |  |  |
| --- | --- | --- | --- |
| 80% | 60% | 20% |  |
| w/o $\omega_{t}$ | 0.2334 $\pm$ 0.0115 | 0.5649 $\pm$ 0.0329 | 3.3820 $\pm$ 0.1704 |
| $\omega_{t}=t$ | 0.2331 $\pm$ 0.0117 | 0.5633 $\pm$ 0.0332 | 3.1881 $\pm$ 0.2170 |
| $\omega_{t}=t^{2}$ | 0.2334 $\pm$ 0.0116 | 0.5647 $\pm$ 0.0333 | 3.3557 $\pm$ 0.1973 |

### G.3 Cross-distribution generalization

**Table 8: Cross-distribution generalization on the Navier-Stokes 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 $\pm$ 0.0162 | 0.2362 $\pm$ 0.0121 | 0.3363 $\pm$ 0.0082 |
| 60% | 0.4990 $\pm$ 0.0260 | 0.5071 $\pm$ 0.0257 | 0.6980 $\pm$ 0.0188 |
| 20% | - | - | 1.9315 $\pm$ 0.0921 |

To evaluate the robustness of our method under distribution shift between training and testing, we conduct experiments where the observation ratio differs between training and inference. Specifically, we investigate whether a model trained on data with a 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 $r_{\text{train}}$ (e.g., 80%), we sample context masks containing a fraction $\alpha$ of the observed points (e.g., 50%), resulting in the model receiving $r_{\text{train}}\times\alpha$ of the total pixels as input (e.g., 40%). At test time with a different observation ratio $r_{\text{test}}$ (e.g., 60%), we maintain the same effective input ratio by sampling $\frac{r_{\text{train}}\times\alpha}{r_{\text{test}}}$ of the available observations as context (e.g., $\frac{40\%}{60\%}=66.7\%$). This strategy ensures the model operates within its learned input distribution while adapting to varying observation densities. Tab. [8](https://arxiv.org/html/2509.20098v1#A7.T8) presents results across different train-test observation ratio combinations, demonstrating that our method maintains reasonable performance even under significant distribution shifts, though performance naturally degrades when the test-time observation ratio substantially deviates from the training distribution.

### G.4 Complete results

We provide the complete experimental results, including standard deviations, to demonstrate the statistical significance and variance of our findings.

**Table 9: Performance comparison on PDE imputation tasks. Each sample represents a temporal sequence of 50 frames, each with $32\times 32$ spatial resolution. Results show the MSE between the reconstructed and the ground truth solutions from the PDE solver, averaged over all timesteps.**
| Method | Shallow Water (feasibility loss, $\times 10^{-8}$) | Advection (simulation MSE, $\times 10^{-1}$) |  |  |  |
| --- | --- | --- | --- | --- | --- |
| 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 $\pm$ 0.1683 | 0.7527 $\pm$ 0.1530 | 0.1030 $\pm$ 0.0004 | 0.1197 $\pm$ 0.0096 |  |
| AmbientDiff | 0.1927 $\pm$ 0.0050 | 0.7504 $\pm$ 0.1119 | 0.1039 $\pm$ 0.0009 | 0.1219 $\pm$ 0.0075 |  |
| Ours | 1 step | 0.1878 $\pm$ 0.0054 | 0.7379 $\pm$ 0.1101 | 0.1035 $\pm$ 0.0008 | 0.1189 $\pm$ 0.0069 |
| 200 steps | 0.1869 $\pm$ 0.0015 | 0.7502 $\pm$ 0.1120 | 0.1037 $\pm$ 0.0009 | 0.1231 $\pm$ 0.0109 |  |

### G.5 Visualization of generated samples

Figure: Figure 4: Imputed results on 2D Shallow Water dataset where 30% of the original data points are available for training.
Refer to caption: figs/sw/sw_10frames_1.png

Figure: Figure 5: Imputed results on 2D Advection dataset where 30% of the original data points are available for training.
Refer to caption: figs/adv/adv_10frames_1.png

Figure: Figure 6: Sample imputation results on 2D Navier-Stokes dataset where 80% of the original data points are available for training.
Refer to caption: x4.png

Figure: Figure 7: Sample imputation results on 2D Navier-Stokes dataset where 60% of the original data points are available for training.
Refer to caption: x6.png

Figure: Figure 8: Sample imputation results on 2D Navier-Stokes dataset where 20% of the original data points are available for training.
Refer to caption: x8.png

Figure: Figure 9: Comparison of original and imputed data from the Navier-Stokes dataset (one missing block). Each sample consists of 100 frames at $64\times 64$ resolution.
Refer to caption: x10.png

Figure: 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 no meaningful generation, confirming our theoretical prediction (Theorem. [1](https://arxiv.org/html/2509.20098v1#Thmtheorem1)) of learning failure in these areas.
Refer to caption: x11.png

Figure: (a) Mean sea level pressure (msl)
Refer to caption: x12.png

## Appendix H LLM usage statement

We used large language models (Claude) to assist with manuscript preparation in the following capacities: (1) improving the clarity and grammatical correctness of our writing through proofreading 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 typographical errors throughout the manuscript.

## Peer Reviews

confidence: 3
contribution: 4
initial_rating: 8
presentation: 3
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?

- 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.
soundness: 3
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.

- 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. 

- 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.

- 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.
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.
The 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).
weaknesses: A bit of nitpicking:
- 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?

- 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.

---

confidence: 4
contribution: 2
initial_rating: 4
presentation: 2
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?

Could 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?

Regarding 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?

What 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?
soundness: 2
strengths: The paper tackles the critical and practical challenge of learning physical dynamics from inherently sparse and incomplete real-world data.

It 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.
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.
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.

Lack 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.

Incomplete 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.

Unaddressed 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.

---

confidence: 3
contribution: 2
initial_rating: 4
presentation: 3
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.
2. 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.
3. 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"?
soundness: 2
strengths: Problem Significance: The paper tackles a crucial and pervasive problem in scientific computing: learning from inherently sparse and incomplete observational data.
Theoretical 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.
Strong Empirical Results: The method demonstrates significant performance improvements across multiple benchmarks, particularly in sparse regimes with very low data coverage (e.g., 1%-20%).
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.
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.
- 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.
- 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.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Q1: What is the term "feature" referring to in the main text? Does it refer to different data modalities or cell grid?
Q2: 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.
Q3: 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.
soundness: 4
strengths: S1: This paper gives a clear literature review and a contrasting comparison with similar models.
S2: The algorithm is backed by rigorous theorems and proofs, and additionally addresses the structured missing prior in reality.
S3: The experiments are thoughtful with ablation studies on the model design and components.
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.
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.
W2: 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.
W3: 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.
W4: The analysis of the paper's limitations is missing in the main text.
W5: The authors claim that some methods are not scalable, but they have not provided a scalability analysis and report on wall-time across baselines.

Minor remarks:
1. Line 723 VAE-based "and" GAN-based methods.
2. The definition of $x_\theta$ is unclear on line 138, and its input and output should be specified.
3. What is the definition of $i$ in eq. 2? Training instance or feature dimension?
4. A cross-reference to the proof of Thm 1 should be included in the paragraph starting from line 172.
5. On line 361, the improvement on MissDiff from noise matching to data matching needs to be quantified.
6. The margin of the caption of Figure 3 is too narrow.


Decision: Accept (Poster)