Controllable diffusion-based generation for multi-channel biological data
Authors: Haoran Zhang
Mingyuan Zhou
Wesley Tansey
Venue: ICLR 2026 Poster (2026)
Keywords: diffusion model
conditional imputation
channel attention
random-masking guidance
imaging mass cytometry
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: Biological profiling technologies, such as imaging mass cytometry (IMC) and spatial transcriptomics (ST), generate multi-channel data with strong spatial alignment and complex inter-channel relationships. 
Modeling such data requires generative frameworks that jointly model spatial structure and inter-channel dependencies and generalize across arbitrary subsets of observed and missing channels. Existing generative models typically assume low-dimensional inputs (e.g., RGB images) and rely on simple conditioning mechanisms that disrupt spatial correspondence and overlook inter-channel dependencies. This work proposes a unified multi-channel diffusion (MCD) framework for controllable generation of structured biological data with complex inter-channel relationships. Our model introduces two key innovations: (1) a hierarchical feature injection mechanism that enables multi-resolution conditioning on spatially aligned observed channels, and (2) two complementary channel attention modules to capture inter-channel relationships and recalibrate latent features. To support flexible conditioning and generalization to arbitrary sets of observed channels, we train the model using a random channel masking strategy, enabling it to reconstruct missing channels given any combination of observed channels as the spatial condition. We demonstrate state-of-the-art performance across both spatial and non-spatial biological data generation tasks, including imputation in spatial proteomics and clinical imaging, as well as gene-to-protein translation in single-cell datasets, and show strong generalizability to unseen conditional configurations.

Published as a conference paper at ICLR 2026
CONTROLLABLE DIFFUSION-BASER GENERATION FOR MULTI-CHANNEL BIOLOGICAL DATA
Haoran Zhang
Mingyuan Zhou
Department of Computer Science
Department of Statistics
University of Texas Austin Austin, TX 78712
University of Texas Austin Austin, TX 78712
hz64530 exas edu
mingyuan, zho uetumcommmbss.o exaxas edd
Wesley Tansey
Computational Oncology
Department of Epidemiology and Biostatistics Memorial Sloan Kettering Cancer Center New York, NY 10065
tanseyw@mskoo.org
ABSTRACT
Biological profiling technologies, such as imaging mass cytometry (IMC) and spa- tial transcriptomics (ST), generate multi-channel data with strong spatial align- ment and complex inter-channel relationships. Modeling such data requires gen- erative frameworks that jointly model spatial structure and inter-channel depen- dencies and generalize across arbitrary subsets of observed and missing chan- nels. Existing generative models typically assume low-dimensional inputs (e.g., RGB images) and rely on simple conditioning mechanisms that disrupt spatial correspondence and overlook inter-channel dependencies. This work proposes a unified multi-channel diffusion (MCD) framework for controllable generation of structured biological data with complex inter-channel reliaannnhi relationships. Our model in- troduces two key innovations: (1) hierarchical feature injection mechanism that enables multi-resolution conditioning on spatially aligned observed channels, and (2) two complementary channel attention modules capture inter-channel rela- tionships and recalibrate latent features. To support flexible conditioning and gen- eralization arbitrary sets of observed channels, we train the model using a ran- dom channel masking strategy, enabling it to reconstruct missing channels given any combination of observed channels the spatial condition. We demonstrate state-of-ther art performance across both spatial and non-spatial biological data generation tasks, including imputation in spatial proteomics and clinical imag- ing, as well as gene-to-protein translation in single-cell datasets, and show strong generalizability to unseen conditional configurations.
INTRODUCTION
Recent advances in generative models enable structured generation, prediction, and imputation across various data domains (Rombach et al., 2022; Zhang al., 2023). Specifically, diffusion models have demonstrated a remarkable capacity to generate high-fidelity samples in natural image and language generation tasks. biology and medical domains, on the other hand, data acqui- sition is often constrained by experimental and clinical limitations. Spatial profiling technologies such as imaging mass cytometry (IMC) (Chang et al., 2017) and sequencing platforms like Xenium (Janesick et al. 2023) are expensive, time-consuming and restricted by physical limitations that only allow measuring a limited number of signals of interest, e.g., around 50 proteins for IMC and 500 to 5000 genes for Xenium. Similarly, in clinical imaging, specific signals may be missing in practice due to patient motion, scan-tin scan-time constraints, or acquisition artifacts (de Verdier et al. 2024). These constraints create pressing need for biological data generation and imputation.

---

Published as a conference paper at ICLR 2026
Unitary correspondence
Complex protein- protein interaction
DNA1
CD206
SMA
SMA
Histone
Ki-67
Original IMC data
DNAZ
CD163
CD31
CD31
DNAZ
DNAT
Most correlated protein
r-0.891
r=0.33
0=0660
00407
7=
r-0.00
ControlNet
f=0.896
/=0,828
0=0.7
7=0.762
F=0.826
r=0.924
Ours
Figure 1: Visual comparison of samples generated by our model with samples generated by Costrolleet-style counterpart. Without the channel attention modules and attention injection mech- anisms, ControlNet-style generative models can handle simple channels with unitary inter-channel correlation, such as DNA1 (the right-most column), but easily fail in those with more complex protein-protein interactions, like tumor markers Ki-67 and aSMA (columns 3-6).
key characteristic of biological data is its multi-channel nature. For instance, in spatial profiling data, each channel is assigned to specific molecule of interest (e.g., proteins n 30 and genes 100), and each pixel (or cell) represents spatially co-registered vector of biologically distinct signals. Unlike natural images with fixed and highly correlated RGB channels, biological channels often exhibit complex inter-channel dependencies. A generative framework that can handle high- dimensional biological data, maintain spatial alignment, accommodate random missing channels, and model diverse inter-channel relationships extends the application of existing technologies and recovers missing information in clinical acquisitions.
In biological data generation, the conditional inputs are typically observed signals of the same sam- ple, which spatially aligned with the target of generation. Such spatial alignment should be preserved throughout the generation process. Naive conditioning mechanisms, such as global em- beddings and flattened concatenation, disrupt spatial correspondence, making them incompatible with the biological data generation objective. On the other hand, the number of channels is large in biological data and their interactions are highly context-dependent. Some channels co-localize only within specific spatial niches or cell types. while others may be mutually exclusive. Therefore, modeling such sparse, nonlinear, and asymmetric dependencies requires generative framework that preserves spatial alignment and adapts to diverse, context-dependent channel semantics in biological data.
In this work, we introduce the multi-channel diffusion (MCD) framework for controllable multi- channel data generation. We demonstrate its performance and generalizability tasks of gener- ating biological profiling data with subsets of observable channels. Our method handles arbitrary combinations of observed and missing channels, maintains spatial alignment, and captures com- plex inter-channel dependencies. To achieve this, propose three key components: (i) amortized conditioning with random channel masking, (ii) adaptive condition injection. and (iii) two comple- mentary channel attention modules. We formalize an amortized conditioning loss and implement
2

---

Published as a conference paper at ICLR 2026
it via random channel masking, enabling single model to generalize across arbitrary conditional configurations of observed and missing channel combinations. Rather than pure element -wiee ad- dition, we reweight latent features in condition injection, allowing the model to dynamically gate spatial conditions. We also combine soft Squeeze-and-Fxcital feature recalibration (used in condition injection) together with transformer-style channel atten- tion module inside UNet blocks to model cross-channel dependencies. Ablation studies confirm that each component improves data generation quality. Together, these components enable our model to generate high-fidelity multi-channel biological data across various condition- -trget combinations within unified diffusion framework. Our model achieves state-of- of-th-------rrr results on both spatial and non-spatial biological data generation tasks and supports test-time controllability for arbitrary conditional configurations.
attention for per-sample global
2 BACKGROUND
2.1 DIFFUSION MODELS
Diffusion models are a class of generative models that approximate the data distribution by mod- eling the gradient of its log-density (Song & Ermon 2019). This is achieved by coupling a forward process that progressively perturbs the data with a reverse process that learns to re- cover the original. The forward process transforms data Xg into a noise distribution XT over T timesteps by injecting noise from known distribution like Gaussian (Sohl-Dickstein al. 2015), q(x, Xo) = Nxx Va,xo. (1 at II) where at is a noise schedule parameterizing the variance at each timestep t. As t approaches T, the data distribution converges standard Gaussian. The reverse process reconstructs data from noise by learning to reverse the cor- ruption introduced during the forward process. This is achieved through a neural network en (), trained to predict the noise added at each timestep (Ho et al. 2020). The denoising process is guided by the objective Ex,.,. .. [||e- cx(x,,,))'" where e ((, IJ the Gaussian noise injected by the forward process. By minimizing this objective, the model learns to reconstruct zo from the noisy a, at any timestep This denoising approach is deeply connected score matching. Specif- ically, the noise prediction co (x(e,))) t) can be used to compute the gradient of the log-density (i.e., so(xx.t)
simple prior, such
score function) as Xx, log q(x,) x vI-a:" Thus, accurately predicting the noise is equivalent to estimating the score function Vx, log (xx), which guides the reverse process, This insight bridges denoising and score matching, making the reverse process a refinement procedure that progressively moves noisy samples back to the data manifold (Song & Ermon 2019).
2.2 STRUCTURED
IMPUTATION
Image imputation is a classical problem in computer vision and generative modeling, where the objective is to recover missing corrupted parts of an image given the observed context. Formally, given an observation c, classical imputation methods aim to estimate the full image E by modeling the conditional distribution p(x c). This problem has been extensively studied in the context of natural images, where x E 3x H HXW and inpainting primarily relies on spatial continuity within the image domain (Chan & Shen 2001).
We generalize this problem to the setting of structured multi-channel data, where each channel corresponds to a semantically distinct signal-such expression. Let e E CRC, C... H denote the observed data, and let x E IR xCC denote the full data, where C C. + Cm 2 1, C, 2 1, and Cmm 2 denote the number of total, observed, and missing channels respectively. Note that when C'm = 0 and C, C = 3, this formulation reduces to the classical RGB image inpainting problem (C = for the grayscale case). When C... = (RGB) and C. 1 (grayscale). it reduces to the classical image colorization problem. This general formulation applies across a broad range of problems at different resolutions. When Hl = W 1, the input reduces to high-dimensional vector, making this formulation applicable to non-spatial data imputation, such as single-cell data, where one can predict single-cell protein expression from single-cell RNA sequencing (scRNA-seg) data. When H, W > 1, the formulation supports spatially structured imaging tasks, such as IMC channel prediction, where both local morphology and global tissue organization contribute the reconstruction target.
spectral band, a protein marker, or gene
3

---

Published as a conference paper at ICLR 2026
b
Noksy Image
Spatial Condition
C
Non-spatial Condition
FILM
Self Attention
Conditional Encader
Channel Attention
Attention
Injection
C
Diffusion Network
Conditional
Network
Diffusion network UNet block
Diffusion network output attention layer
Condittonal network UNet block
E
Conditional convolution layer
Figure 2: Overview of the proposed MCD for multi-channel data generation. (a) Custom chan- nel attention: UNet block with channel attention module and attention injection block that model the inter-channel relationships. (b) Hierarchical feature injection: parallel conditional network encodes the spatial condition, producing features at different resolutions that are injected into the corresponding UNet block in the diffusion network.
METHODS
We propose diffusion- based generative framework, MCD, for multi-channel biological data that conditions on arbitrary subsets of observable channels. Let x E RC* denote the full-channel data and x an observed subset, where x and are spatially aligned. The goal is to learn the conditional distribution p(x[c) in way that respects the spatial structure of the data while remaining flexible any condition c. This requires resolving several intertwined challenges: 1) the generated targets should be spatially aligned with the conditions; 2) the conditions should be modeled as structured, multi-resolltion information; 3) correlations among biological channels are complex; and 4) the model must generalize across arbitrary condition-targe combinations at test time. To address these challenges, MCD employs a dual-network architecture consisting of diffu- sion network that denoises the noisy target X, and conditional network that encodes the observed channels c (Figure 2p). At each resolution level f, the diffusion encoder produces features Dz(x;) while the contextual encoder provides aligned features Ec(c). These conditional representations hierarchically injected into the diffusion network at corresponding resolution to ensure alignment and effective spatial conditioning. Channel attention modules capture inter-channel dependencies, and random channel masking training ensures generalizability across channels.
3.1 SPATIALLY ALIGNED FEATURE INJECTION FOR STRUCTURED CONDITIONING
Generating realistic multi-channel biological data requires preserving spatial correspondence with the observed channels and incorporating the conditioning signal that guides both local details and global structure. To achieve this, we introduce a hierarchical feature injection mechanism. At each resolution level f, contextual features Eg(c) are spatially aligned with and injected into the diffusion features (z,);) through
Zr = Dx(x)) + SB(EE(C)
allowing for the flexible injection of spatial conditions ("Attention Injection" block in Figure 2 The SE(-) is the sineezze-and-Exiitttte block that serves soft channel attention mechanism (see section 3.3 for detail) and selectively injects the condition while the spatial alignment is pre- served. SE Ee(c)) maintains spatial dimensions and applies channel-wise gating to the condition
4

---

Published as a conference paper at ICLR 2026
feature map before dlement-wise addition. This design allows the model to condition not just fixed global representation с but on a series of contextual features {Ee(c)) that vary across resolution. This setup reflects the intuition that certain pattems in e.g., global motifs versus lo- cal structures, may depend on different aspects of c. Early encoder layers focus more on the local structures, while late layers draw high-level, global structures.
3.2 TRAINING WITH RANDOM CHANNEL MASKING
Different experiments may measure different sets of signals. Rather than training separate models to predict individual channels, our goal is to produce full multi-channel outputs (i.e., the complete protein panel). Constructing our model to condition on any subset enables us to train it across multiple datasets and flexibly apply
each. To create a unified model, propose random channel masking strategy during training (Algorithm D. In each iteration, we randomly sample subset of channels S, C (I....,C C) observed, and mask the remaining Sm = {_..... C)S.. The masked subset Xs.. is zeroed out in spatial condition c, and the model reconstructs the full panel RC* We optimize the standard EDM objective (Karras al. 2022) on the full-channel target x, and the binary mask is only applied to the conditions.
Algorithm Random Channel Masking
Require: Full data RCC x masking probability p
1: for each training iteration do
2: Sample observed S, (l,..., CJ, where lies, Bern(p)
3: Construct condition e such that
X i S.
G: 0,
otherwise
4: Diffusion step with target x and condition c
5: end for
Randomly varying S. during training encourages the model to learn conditional generation under diverse partial contexts. As a result, the trained model can generalize to arbitrary combinations of observed and missing channels at test time, including unseen configurations. At the same time, because the model always reconstructs all channels, it avoids the need for channel-specifie heads or separate training for each channel, enabling comprehensive downstream analysis and one-stop training.
This training procedure implicitly defines an amortized conditional learning objective:
Pee)lEe |||e eg (Xe, &, e)
(2)
where p(c), with E C, denotes the distribution of conditional configuration (i.e., combinations of observable channels) over the conditional space C. Minimizing eq. encourages the model to learn single amortized estimator Cp (x,x,,,, t, such that (, R log p(x for arbitrary condition Under standard diffusion assumptions, random channel masking ensures that the model trained on random conditional subsets, and the learned model generalizes across configurations and implicitly models the family of conditional distributions {p(z|c), C). Therefore, this for- mulation frames random masking as an amortized inference over the condition space. Although not formally addressed, prior work has demonstrated its effectiveness across various domains (Gersh-) man & Goodman, 2014) Marino et al. 2018). We provide theoretical justification in Appendix A Empirically. we observe that this strategy enables robust generalization to unseen conditional con- figurations at test time.
3.3 CHANNEL ATTENTION STRUCTURED FEATURE MODULATION
With random channel masking, model aims to handle arbitrary conditional configurations, where the observed channels may vary across samples and tasks. Since all conditions c are encoded with the same conditional network, the model needs to select and reweight latent features adaptively
5

---

Published as a conference paper at ICLR 2026
depending the conditions. To enable dynamic feature recalibration in the latent space, intro- duce two complementary channel attention modules with distinct designs, each focusing on specific modeling goals.
The first is a lightweight attention mechanism ("Attention Injection" block in Figure 2 inspired by the Squeeze-ande and-Excitation (SE) network (Hu et al. 2018). Given latent feature map DxDxXW H...
we compute
a GAP(z)
W (W2 <WW, a))
z' = W
where á is a non-linear activation function (e.g., ReLU), GAP stands for global average pooling, and a is the sigmoid function. This approach scales each latent channel by a learned weight conditioned on the global context and enables per-channel feature reweighting that aligns with the dynamic feature selection requirement for multi-scale conditional injection (Section [.I). This SE-based attention mechanism is lightweight and stabilizes hierarchical feature injection.
On the other hand, we also perform a self-attention over all latent channels ("Channel Attention"
blocks in Figure 2
Xstar R
ERD
= W K Xua WK,
(4)
Q = Xna WQ:
V ==xaaW We
QKIT T
= softmax
Xnat
AV
(5)
Vd
Compared to SE channel attention, this design is more expressive and can capture higher-order de- pendencies among latent channels. This feature makes the transformer-based channel attention par- ticularly useful in UNet blocks, where the model infers missing information across latent channels. Rather than reweighting channels independently. it allows information to propagate across chan- nels through learned interactions. This is particularly important in biological data, where channel relationships are often nonlinear, asymmetric, and context-dependent.
MCD includes an additional soft channel attention block, in the same way as Equation at the final stage of the model, which maps the latent channels to data channels. Let the final latent representa- tion before projection be E R Dx where D is the number of latent channels. The standard diffusion model produces y E xC with a single output layer on z. To model cross-channel dependencies, we add an additional layer that computes
attn y Conv (SE(y))
While latent-space attention captures dependencies among hidden features, the output-space atten- tion provides a final channel- recalibration in the data space.
Together, these mechanisms provide both adaptive gating (SE) and structured channel interaction modeling (self-attention), enabling robust feature modulation across heterogeneous conditional con- figurations and improving performance on tasks involving complex, structured channel relation- ships.
4 RELATED WORK
Image inpainting with conditional diffusion models. Recent work has applied diffusion- generative models to image inpainting and completion tasks (Song Ermon 2019 Saharia al., 2022). These models primarily work on natural and grayscale images, where the number of channels limited (n S 3), and the goal is to reconstruct spatially masked regions based on the surrounding context. These models often take conditionals like class labels (Dhariwal & Nichol, 2021), text embeddings (Ramesh et al. 2022 Nichol et al. (002),, or segmentation maps (Rombach et al. 2022). These conditionals are typically injected via concatenation at the input or embedding injection via FiLM modulation. This approach ignores the potential spatial alignment between the conditionals and the generative target. More recent approaches like ControlNet (Zhang et al. 2023) and BrushNet (Ju et al.) 2024) introduce multiscale conditioning mechanisms that are spatially aligned and applied post hoc pre-trained Stable Diffusion models. While these methods preserve spatial alignment,

---

Published as a conference paper at ICLR 2026
they assume low-dimensional inputs and do not address the challenges posed by high high-dimensional structured data with intricate inter-channel dependencies. Moreover, because their conditioning modules trained separately from the core generative model, they lack end-to-end coordination between condition encoding and generation.
Dynamic guidance for conditional diffusion. Classifier-free guidance (CFG) (Ho Salimans, 2021) introduces a flexible training scheme for conditional diffusion models by randomly dropping conditioning inputs and jointly training the model on both conditional and unconditional objectives. While originally developed for low-dimensional conditioning signals such as class labels, the core idea can be generalized: using input masking during training to enable flexible guidance at test guidance, where the spatial conditions
time. We adopt this principle in the form of are randomly masked during training. Unlike CFG, which interpolates between conditional and unconditional predictions at inference time, method samples conditioning subsets during training to learn a single unified conditional model. Specifically, the model observes a random subset of input channels and is trained reconstruct the full panel. This dynamic masking encourages the model to generalize across different conditionals, making it suitable for tasks where the available condition varies across samples. This approach enables multi-channel prediction and supports generation under arbitrary conditioning subsets without retraining or architectural changes.
Channel attention mechanisms. While channel-wise attention has been explored in vision ar- chitectures like SENet (Hu et al. 2018), most diffusion- based models focus spatial attention and overlook channel-wise relationships. However, this becomes a significant limitation in multi- channel biological data. A few recent works in diffusion-based colorization have begun explore this direction: FCNet (Zhu et al. 2024) introduces spatially decoupled color representations tailored to facial regions, and ColorPeel (Butt al. 2024) learns disentangled geometric shapes and color embeddings in latent space. These models can be thought of as coarse region-level attention over RGB channels and designed for tasks involving dozens hundreds of semantically distinct channels.
5 EXPERIMENTS
5.1 SINGLE-CELL MODALITY PREDICTION
To evaluate MCD's performance in non-spatial biological data, we begin with the standard CITE-seq benchmark of predicting protein expression from paired scRNA-seg. where the targets single- cell protein expression, and the conditions are scRNA-seg expression. The model is trained without random channel masking since both the target and the condition channels are fixed. We compare MCD to a suite of state-of-the-al methods: Kernel Ridge Regression (Guanlab-dengkw (Lance al. 2022), MultiVI (Ashuach al.] 2023). GLUE (Cao Gao 2022), ScMM (Minoura al.] 2021 and UnitedNet (Tang et al. Across four datasets: peripheral blood mononuclear cells (PBMC, Hao et al. (2021)) cord blood mononuclear cells (CBMC, [Stoeckius al. (2017)) bone marrow mononuclear cells (BMMC, Lance et al. (2022)). and hematopoietic stem and progenitor cells (HSPC,
(HSPC,,Nestorowa /Nstorowa et al. (2016)) our model consistently outperforms the existing single-cell
modality translation methods (Table D. Performance is reported as average Pearson correlation between predicted and measured expression levels across proteins (r,) and cells (re). Importantly, framework achieves the highest protein-level correlation (rp), the more biologically relevant metric, in every dataset. This reflects the model's ability to capture gene-protein relationships at scale.
We also distilled our diffusion model with SiD (Zhou et al. 2023 into a one-step generation vari- ant, which retains comparable accuracy while reducing inference cost by two orders of magnitude. This suggests that MCD not only yields high-accuracy generation but also ready for large-scale deployment.
7

---

Published as a conference paper at ICLR 2026
Table 1: Benchmarking multimodal diffusion approach against existing modality prediction methods on the gene-to-protein prediction task on four (PBMC, CBMC, BMMC, and HSPCs) datasets, te shows the cell-wise correlation with the ground truth, and Tp shows the protein-level correlation.
Method
PBMC
CBMC
BMMC
HSPC
Ve
Vp
Ve
Yp
Te
Yp
Te
Jp
KRR (Guanlab-dengkww)
0.908
0.646
0.863
0.006
0.870
0.094
0.820
0.059
MultiVI
0.088
0.069
0.103
0.032
0.082
0.054
0.045
0.035
GLUE
0.659
-0.004
0.623
0.054
0.589
0.012
0.549
0.024
UnitedNet
0.870
0.518
0.377
0.628
0.625
0.634
0.310
0,436
scMM
0.793
0.521
0.736
0.517
0.853
0.625
0.724
0.598
Ours (500 steps)
0.880
0.673
0.962
0.763
0.879
0.685
0.865
0.647
Ours + SiD (1 step)
0.874
0.672
0.962
0.759
0.875
0.682
0.865
0.642
5.2
SPATIAL DATA IMPUTATION
5.2. SINGLE DATASET IMPUTATION
To evaluate the model's 's performance spatial biological data, we apply it to spatial proteomics data from imaging mass cytometry (IMC), where each pixel corresponds to spatially co-registered multi-protein expression signals. We evaluate our method on two IMC datasets: a lung cancer co- hort (Yoffe et al. 2025) (8 patients) and breast cancer cohort (Jackson et al.) 2020) (6 patients), with 50 and 35 co-registered protein channels, respectively. The spatial data are normalized per channel and partitioned into 64 x 64 patches, which are further downsampled to 16 16 by pre- trained encoder. For each dataset, we compare our generation to two data-specific baselines: the protein with the highest spatial correlation and kernel ridge regression, as well as structured condi- tional diffusion models, ControlNet (Zhang et al.] 2023). We also include domain-specific methods, STEM (Zhu al., 2025) and Virtues (Wenckstern et al. 2025). We evaluate under two protocols: (1) single-channel, where a dedicated model is trained per target channel, and (2) multi- chlti--hhannl,, where a single model is trained with random channel masking. In the single-channel prediction setup, we mask one protein channel at time as the target and use the remaining channels as the spatial condition. On the other hand, in the multi-channel prediction setup, the observed channels are randomly selected by independent (0.9( random variables, and the remaining channels are set to zero. The model is trained to reconstruct the full-channel data. To explicitly evaluate gen- eralization across unseen conditional configurations, includes channel subsets not observed during training. Successful reconstruction under previously unseen masking patterns provides evidence for true amortized generalization across the conditional space rather than memorization of fixed condition configurations.
ensure that the masking used at test time
method outperforms all existing approaches on both IMC datasets. Most baseline models fail to outperform the best linear predictor, as well as the most correlated individual ob- served protein. By contrast, our method consistently outperforms the baselines across each protein channel. Fixing the missing channel during training improves performance in the single-channel pre- diction mode compared the multi-channel setup, which reflects the capacity trade-off in amortized multi-task learning. When the target channel is fixed during training, the model can dedicate its ca- pacity to modeling a single conditional distribution. contrast, the multi-channel model amortizes its capacity across all possible missing-channe configurations. Despite this broader objective, the multi-channel model still outperforms all baselines, demonstrating that the random masking strategy enables strong generalization with high predictive accuracy. These results validate the effectiveness of random- random-masking guidance and support practical use cases where multi-channel prediction is re- quired. Visualization of the generated IMC samples (Figure (), ablation studies (Appendix B.2) and quantifications of uncertainty and correlations between generated channels (Appendix B.3 are provided in the appendix.
Table
8

---

Published as a conference paper at ICLR 2026
Table 2: Comparison of predictive correlation across different methods in spatial prediction tasks. Our method outperforms the diffusion, single-pre iigle-protein, and linear baselines, and diffusion based ControlNet, domain-specific models Stem and Virtues.
Method
Breast
Lung
Most correlated protein Kernel ridge regression
0.481
0.506
0.489
0.527
Virtues Wenckstern et al. 2025
0.398
0,425
Stem (Zuu et al. (0255
0.403
0.475
NontrolNet (Zhng et al. 2023
0.452
0.537
Ours (single channel)
0.667
0.703
Ours (multi channel)
0.596
0.647
1.2
MLA
CD208,
DNK1
Generative corr. w/ union set
C268 H33FaPP CD4
Generative corre ation
PanCK
#CD16
CDGB
Granaymes>
cD3
NER1
CD14
am
0.
MNP9, GITR COEMM
MIA Dis
60577
CD8y
CD56
20.1
CD11b, GranzymeB CD33
aSMA
CDe
CD31
Fanck L1
CD20
TOX
SFTPC
C33
2110
JSMA
CD8a
62
03
04
24
07
Protein-wise best correlation
Generative corr w/ intersection set
Figure 3: (a) Benchmarking generative performance (Pearson's r) against the most similar protein on the lung cancer dataset. (b) Cross-dataset generalization study results. Using the union of all protein channels, those missing from the test dataset, consistently outperforms using only the overlapping set of protein channels.
5.2.2 MULTI-DATASET GENERALIZATION
We further test model's ability to learn from multiple datasets with partially overlapped protein channels. Specifically, we combine the lung and breast cancer IMC datasets. The two datasets share 23 proteins common. We evaluate our model under two training schemes: (I) an intersection setting, where we keep only the proteins observed in both datasets, and (2) a union setting, where include the union of all protein channels and zero-pad missing channels in each dataset. Both setups share the same model architecture and training strategy with random masking, and predict the intersection set test time.
As shown in FigureBp. the union setup consistently outperforms the intersection counterpart, achiev- ing higher average Pearson correlation in sample generation. The union setup has broader cov- erage and enables the model to learn a richer set of inter-channel dependencies and improves its robustness to missing data. This result empirically validates that random channel masking enables principled joint learning under heterogeneous supervision. By zero-padding unobserved channels and sampling conditional subsets during training, the model does not require perfect channel align- ment across datasets. Instead, it leams infer conditional structure from partially overlapping panels. The performance in the union setting demonstrates MCD and random channel masking as principled method for malti-dataset integration with empirical advantages.

---

Published as a conference paper at ICLR 2026
5.3 MRI MODALITY SYNTHESIS
To test the generalizability of our framework, we evaluate on the BraTS benchmark (de Verdier et al., 2024 for missing-modalie magnetic resonance imaging (MRT) synthesis, where the goal is construct unacquired MR sequences given observed ones. This task is clinically relevant in brain tumor imaging, as missing scans are common practice due to acquisition time constraints and -atient-specific factors. We follow the BraTS guidelines, using tissue structural similarity (SSIM) for image fidelity and Dice coefficient (DICE) for downstream tumor segmentation quality as eval- uation metrics.
Table summarizes results across SOTA baselines, including Pix2pix (Isola al. 2017). HF-GAN (Cho et al. 2024), and SwinUNETR (Pang et al. | (025), where the latter two are the winner and the runner-up of the BraTSs 2024 challenge, respectively. Our method achieves the highest scores across both SSIM and DICE, improving both the structural similarity and clinical utility. Gains in DICE demonstrate that our approach not only reconstructs visually faithful images but also preserves tumor structures critical for segmentation. Compared with HF-GAN and SwinUNETR, which represent state-of-the- GAN- and transformer- baselines, our model shows consistent improvements in SSIM and DICE. Visualization of the generated MRI samples is provided in the appendix (Figure 7
Table 3: Comparison of segmentation accuracy and structural similarity across different synthesis methods with DICE and SSIM scores across tumor, tissue, and global regions.
Method
DICE f
SSIM t SSIMheathy t
SSIM t
SSIM glohal
pix2pix
0.549
0.719
0.570
0.583
0.807
HF-GAN
0.714
0.761
0.604
0.615
0.919
SwinUNETR
0.709
0.759
0.628
0.637
0.916
MCD
0.738
0.774
0.631
0.643
0.928
CONCLUSION
We introduce MCD as a diffusion-basee generative model for the controllable generation of multi- channel biological data, which is generalizable to arbitrary combinations of observed and missing channels. By combining amortized conditioning via random channel masking, hierarchically aligned feature injection, and structured channel attention, MCD captures complex inter-channel depen- dencies while maintaining spatial alignment. Our channel masking and simulation-basee analyses highlight the effectiveness of MCD, and our experiments MRI modality synthesis demonstrate that the method can provide tangible benefits in real-world applications, such as recovering missing clinical signals to support downstream analyses.
MCD offers an in silico extension to the experimentally constrained profiling panel, and recovers corrupted or incomplete acquisitions. Together, these properties suggest that this framework can serve as a scalable backbone for structured biological data modeling. In this sense, MCD provides a step toward foundation model for spatial and multi-modal biological profiling, where single gen- erative prior supports imputation, synthesis, and multi-dataset integration. While this work focuses on methodological development and demonstrates the usefulness of the proposed framework in bi- ological image generation tasks, future directions include scaling to larger and more diverse spatial cohorts, incorporating richer biological priors, and facilitating biologically meaningful discovery.
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
The complete theoretical justification and assumption needed are provided in Appendix A Code is available at https//giihubb.comm/maeeel//AAAl
REFERENCES
Tal Ashuach, Mariano I. Gabitto, Rohan V. Koodli, Ciuseppe-Antonio Saldi, Michael I. Jordan, and Nir Yosef. MultiVI: deep generative model for the integration of multimodal data. Nature Methods, 2023.
Peter L. Bartlett and Shahar Mendelson. Rademacher and gaussian complexities: Risk bounds and structural results. Journal of Machine Learning Research, 3(Nov):463 -822 2002.
Muhammad Atif Butt, Kai Wang, Javier Vazquez-Corral, and Joost van de Weijer. ColorPeel: Color
prompt learning with diffusion models via color and shape disentanglement. Proceedings of European Conference Computer Vision (ECCV), 2024.
Zhi-Jin Jie Cao and Ge Gao. Multi-omics single-cell data integration and regulatory inference with graph-linken embedding. Nature Biotechnology, 40:1458- 1466, 2022.
Tony F. Chan and Jianhong Shen. Nontexture inpainting by curvature-driven diffusions. Journal of visual communication and image representation, 12(4):436-449, 2001.
Qing Chang, Olga Ornatsky, Iram Siddiqui, Alexander Loboda, Vladimir Baranov, and David W Hedley. Imaging mass cytometry. Cytometry Part A, 91(2):160-169, 2017.
Jihoon Cho, Seunghyuck Park, and Jinah Park. Two-stage approach for brain MR image synthesis: 2d image synthesis and 3D) refinement. arXiv preprint (.Ni::2400.102663 2024.
Corinna Cortes, Spencer Greenberg, and Mehryar Mohri. Relative deviation learning bounds and generalization with unbounded loss functions. Annals of Mathematics and Artificial Intelligence, 85(1):45-70, 2019.
Maria Correia de Verdier, Rachit Saluja, Louis Gagnon, Dominic LaBella, Ujjwall Baid, Nourel Hoda Tahon, Martha Foltyn-Dumitru, Jikai Zhang, Maram Alafif, Saif Baig, Ken Chang, Gennaro D' Anna, Lisa Deptula, Diviya Gupta, Muhammad Ammar Haider, Ali Hussain, Michael Iv, Marinos Kontzialis, Paul Manning, Farzan Moodi, Teresa Nunes, Aaron Simon, Nico Soll- mann, David Vu, Maruf Adewole, Jake Albrecht, Udunna Anazodo, Rongrong Chai, Verena Chung. Shahriar Faghani, Keyvan Farahani, Anahita Fathi Kazerooni, Eugenio Iglesias, Flo- rian Kofler, Hongwei Li, Marius George Linguraru, Bjoern Menze, Ahmed W. Moawad, Yury Velichko, Benedikt Wiestler, Talissa Altes, Patil Basavasagar, Martin Bendszus, Gianluca Brug- nara, Jaeyoung Cho. Yaseen Dhemesh, Brandon K. K. Fields, Filip Garrett, Jaime Gass, Lubomir Hadjiiski, Jona Hattangadi-Gluthh Christopher Hess, Jessica L. Houk, Edvin Isufi, Lester J. Layfield, George Mastorakos, John Mongan, Pierre Nedelee, Uyen Nguyen, Sebastian Oliva, Matthew W. Pease, Aditya Rastogi. Jason Sinclair, Robert X. Smith. Leo Sugrue, Jonathan Thacker, Igor Vidic, Javier Villarueva-Meyer, Nathan S. White, Mariam Aboian, Gian Marco Conte, Anders Dale, Mert R. Sabuncu, Tyler M. Seibert, Brent Weinberg, Aly Abayazeed, Ray- mond Huang, Sevcan Turk, Andreas M. Rauschecker, Nikdokht Farid, Philipp Vollmuth, Ay- man Nada, Spyridon Bakas, Evan Calabrese, and Jeffrey D. Rudie. The 2024 brain tumor segmentation (BraTS) challenge: Glioma segmentation on post-treatment MRI. arXiv preprint arXiv:2405.18.80 18368, 2024.
Prafulla Dhariwal and Alex Nichol. Diffusion models beat gans image synthesis. In Advances in 2021.
Neural Information Processing Systems, volume 34, pp.
Samuel Gershman and Noah Goodman. Amortized inference in probabilistic reasoning. In Pro- ceedings of the Annual Meeting of the Cognitive Science Society, 2014.
Yuhan Hao, Stephanie Hao, Erica Andersen-Nissen. William M. Mauck, Shiwei Zheng, Andrew Butler, Maddie Lee, Aaron J. Wilk, Charlotte Darby, Michael Zager, Paul Hoffman, Marlon Stoeckius, Efthymia Papalexi, Eleni P. Mimitou, Jaison Jain, Avi Srivastava, Tim Stuart, Lamar M.

---

Published as a conference paper at ICLR 2026
Fleming, Bertrand Yeung, Angela J. Rogers, Juliana M. McElrath, Catherine A. Blish, Raphael Gottardo, Peter Smibert, and Rahul Satija. Integrated analysis of multimodal single-cell data. Cell, 184(13):3573-3587, 2021.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. In NewrlPS 2021 Workshop on Deep Generative Models and Downstream Applications, 2021.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, 2020.
Jie Hu, Li Shen, and Gang Sun. Squeeze-and-excitan networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 713 7132-7141, 2018.
Phillip Isola, Jun-Yan Zhu, Tinghui Zhou, and Alexei A Efros. Image-to-image translation with conditional adversarial networks. In Proceedings of the IEEE conference on computer vision and pattern recognition (CVPR), 2017.
Hartland W. Jackson, Jana R. Fischer, Vito R. т. Zanotelli, H. Raza Ali, Robert Mechera, Savas D. Soysal, Holger Moch, Simone Muenst, Zsuzsanna Varga, Walter Weber, and Bernd Boden- miller. The single-cell pathology landscape of breast cancer. Nature, 578(7796):615 6-00000066666600. 2020.
Amanda Janesick, Robert Shelansky, Andrew D. Gottschoo Florian Wagner, Stephen R. Williams, Morgane Rouault, Ghezal Beliakoff, Carolyn A. Morrison, Michelli F. Oliveira, Jordan T. Sicher- man, Andrew Kohlway, Jawad Abousoud, Tingsheng Yu Drennon, Seayar H. Mohabbat, Sarah E. B. Taylor, and 10x Development Teams. High resolution mapping of the tumor microenvi- ronment using integrated single-cell, spatial and situ analysis. Nature Communications, 14(1): 8353, 2023.
Xuan Ju, Xian Liu, Xintao Wang, Yuxuan Bian, Ying Shan, and Qiang Xu. BrushNet: A plug-and and- play image inpainting model with decomposed dual-branch diffusion. In European Conference on Computer Vision, 2024.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Advances in Neural Information Processing Systems, 2022.
Christopher I Lance, Malte D. Luecken, Daniel B. Burkhardt, Robrecht Cannoodt, Pia Rautenstrauch, Anna Laddach, Aidyn Ubingazhibov, Zhi-Jie Cao, Kaiwen Deng, Sumeer Khan, Qiao Liu, Niko- lay Russkikh, Gleb Ryazantsev, Uwe Ohler, NeurIPS 2021 Multimodal data integration compe- tition participants, Angela Oliveira Pisco, Jonathan Bloom. Smita Krishnaswamy, and Fabian J. Theis. Multimodal single cell data integration challenge: Results and lessons learned. In Pro- ceedings the NeurIPS 2021 Competitions and Demonstrations Track, volume 176, pp. 162-176, 2022.
Joe Marino, Yisong Yue, and Stephan Mandt. Iterative amortized inference. In International Con- ference on Machine Learning, pp. 3403-3412. PMLR, 2018.
Kodai Minoura, Ko Abe, Hyunha Nam, Hiroyoshi Nishikawa, and Teppei Shimamura. A mixture- of-experts deep generative model for integrated analysis of single-cell multiomics data. Cell Reports Methods, 1050:00077,, 2021.
Sonia Nestorowa, Fiona K Hamey, Blanca Pijuan Sala, Evangelia Diamanti, Mairi Shepherd, Elisa Laurenti, Nicola K Wilson, David G Kent, and Berthold Göttgens. A single-cell resolution map of mouse hematopoietic stem and progenitor cell differentiation. Blood, The Journal of the American Society of Hematology, e2, 2016.
Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranay Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen. GLIDE: Towards photorealistic image generation and editing with text-guided diffusion models. In Proceedings of International Conference on Machine Learning (ICML). 2022.
Haowen Pang, Weiyan Guo, and Chuyang Ye. Multi-modal brain MRI synthesis based on Swin- UNETR. arXiv preprint arXiv:2506.02467, 2025.
12

---

Published as a conference paper at ICLR 2026
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with clip latents. arXiv preprint arXiv:2204.06125 2022.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Con- ference on Computer Vision and Pattern Recognition, pp. 10684-10695, 2022.
Chitwan Saharia, William Chan, Huiwen Chang, Chris A. Lee, Jonathan Ho, Tim Salimans, David Fleet, and Mohammad Norouzi. Palette: maage-to-image diffusion models. In SIGGRAPH, 2022.
Jascha Sohh-Diccstein, Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamies, In Proceedings of the 32nd International Conference on Machine Learning (ICML), volume 37, pp.
2015.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution.
In Advances Neural Information Processing Systems (NeurIPS), volume 33, pp. 11895-11907, 2019.
Marlon Stoeckius, Christoph Hafemeister, William Stephenson, Brian Houck-Loomis, Pratip K Chattopadhyay, Harold Swerdlow, Rahul Satija, and Peter Smibert. Simultaneous epitope and transcriptome measurement in single cells. Nature Methods, 14(9):865-866 160,8665-88, 2017.
Xin Tang, Jiawei Zhang, Yichun He, Xinhe Zhang. Zuwan Lin, Sebastian Partarrieu, Emma Bou Hanna, Zhaolin Ren, Hao Shen, Yuhong Yang. Xiao Wang, Na Li, Jie Ding, and Jia Liu. Explain- able multi-task learning for multi-modality biological data analysis. Nature Communications, 2024.
Maya Varma, Ashwin Kumar, Rogier Van der Sluijs, Sophie Ostmeier, Louis Blankemeier, Pierre Joseph Marcel Chambon, Christian Bluethgen, Jip Prince, Curtis Langlotz, and Akshay $ Chaud- hari. Medvae: Efficient automated interpretation of medical images with large-scale generalizable autoencoders. In Medical Imaging with Deep Learning, 2025.
Johann Wenekstem, Eeshaan Jain, Kiril Vasilev, Matteo Pariset, Andreas Wicki, Gabriele Gut, and Charlotte Bunne. Al-powered virtual tissues from spatial proteomies for clinical diagnostics and biomedical discovery. arXiv preprint arXiv:2501.0603: 2025.
Liron Yoffe, Bhavneet Bhinder, Sung Wook Kang, Haoran Zhang, Arshdeep Singh, Hiranmayi Ravichandran, Geoffrey Markowitz, Mitchell Martin, Junbum Kim, Chen Zhang, Olivier Ele- mento, Wesley Tansey, Stewart Bates, Timothy E. McGraw, Alain Borczuk, Hyun-Sung Lee, Nasser K. Altorki, and Vivek Mittal. Acquisition of discrete immune suppressive barriers con- tributes the initiation and progression of preinvasive to invasive human lung cancer. bioRxiv, pp. 2024-12, 2025.
Lymin Zhang, Anyi Rao, and Maneesh Agrawala. Adding conditional control to text-to-image diffusion models. In International Conference Computer Vision (ICCV), 2023.
Mingyuan Zhou, Huangjie Zheng. Zhendong Wang. Mingzhang Yin, and Hai Huang, Score identity distillation: Exponentially fast distillation of pretrained diffusion models for one-step generation. In Proceedings of the 41st International Conference on Machine Learning. 2023.
Hangyan Zhu, Ming Liu, Chao Zhou, Zifei Yan, Kuanquan Wang, and Wangmeng Zuo. Learn- ing spatially decoupled color representations for facial image colorization. arXiv preprint arXiv:2304.122 12345, 2024.
Sichen Zhu, Yuchen Zhu, Molei Tao, and Peng Qiu. Diffusion generative modeling for spatially re- solved gene expression inference from histology images. In International Conference on Learning Representations (ICLR), 2025.
13

---

Published as a conference paper at ICLR 2026
A THEORETICAL JUSTIFICATION OF RANDOM MASKING
We provide a theoretical justification for our random masking training strategy. In appendix A.1 we show that the batch-wise empirical amortized loss upper bounds the true amortized loss. In appendix A.2] we demonstrate that the model's performance in an arbitrary condition configuration controlled, with high probability, under the empirical assumption variance. We define amortized loss
L = L, [((c, t, c))l
ar [||e c c)
and the empirical loss
N [||c (,
iN
(8)
N
under the bounded loss assumption. Assumption 1. (Bounded loss) The loss ((s.e.r(e,,(((xxt.t).))) : D D + [0, d] asymptotically for trained co by minimizing LN-
Since (0 I) has an expected squared norm proportional to its dimension, the per-sample squared loss is bounded in expectation. Specifically, for normalized input and standard Gaussian noise, the loss is upper bounded by the variance of i... [[[[<|]]] D == CHW. and significantly lower as the model converges. Therefore, for the following justifications, we assume the loss lies within fixed scale d S D asymptotically, and will provide empirical evidence in appendix B.T
A.1 EMPIRICAL BOUND OF THE AMORTIZED LOSS
Theorem 1. [Bartlett and Mendelson '02] Consider loss function L Y A [0, 1] and dominating cost function o Y A [0, 1]. Let F be a class of functions mapping from to A and let (Xi, Y:) be independently selected according the probability measure P. Then, for any integer n and any 0 < < 1, with probability at least 1 - over samples of length nı, every f in F satisfies
8 (2(/2
E [CC, [()) R.,, (oY, R2 F)
n
where 6o F (z. oly f(x)) (l. E F}, and R., is the Rademacher complexity.
Theorem follows McDiarmid's inequality, and the detailed proof is provided by (Bartlett & Mendelson 2002).
Corollary 1. Under the bounded-loss assumption, with probability at least 1 - 6, is upper-bounded by LN up to an additive constant.
Proof. Define function class Fe (co(xz,t,c)) for the trainable diffusion model. Let L (c. fe) En.mo data space such that x, e D, and eq. (9) becomes
x D + [0, 1], where Enom = 3 is the normalized squared loss and D the
d
N Ex [C (e, fc, (Ze, t))] + (from Fb) + dy
8ln(2/8)
&.,.EE, [foom (e, fo(xe, t))] <
N
N
1
N
88ln(2/))
E,Er, (e, t)]]
Er [dl (e, (Ie, t))] + Ry (dl Fe) dy
N
N
1
8 8ln(2/8)
E.E. [ele, t))]
E, E [ec, (x (e Fo) + dy
N
L(Ee) & (e( + C(8)
(10)
14

---

Published as a conference paper at ICLR 2026
where C()) Ry(eoF)+ 8l1(2/0), By Bartlett and Mendelson Bartlett & Mendelson 2002 Ry (dFo) = (Fo) and (f Fa) S Rn where & Lipschitz continuous with constant
Le. Under the bounded-loss assumption, reasonable loss functions, including the squared loss
use, satisfy the Lipschitz condition on [0, d). and C()) S L,Rx (Fe) + V a determined by the model architecture and ((5) is independent of f once the function class Fa is fixed. Empirically,
8 $ln(2/0), where Ry (Fe)
find d < 1 shown in fig. Therefore, the true amortized loss is upper-bounded by the empirical loss up to an additive constant.
Similarly, one can get the lower bound of same fashion
(e(()) - +
(11)
and similar results can also be drawn from Corollary 18 from the work by Cortes et al. (2019) Therefore, we show that minimizing our empirical loss also minimizes the true amortized loss.
A.2 CONDITION-SPECHIS BOUND BY AMORTIZED LOSS
Lemma 1. Assuming relatively low variance, i.e. S ps, with high probability, the per-condition loss is of the same magnitude as the amortized loss μ4, i.e. Er,zo [e(e, folWe, f))] = O(u)
Proof. Define the condition-specifii loss
Le = Eo Solx,, t))]
(12)
By the one-sided Chebyshev inequality.
σ²
Pr(Le BoLLo] 2 a) S
σ2 + a²
(13)
Pr(Le 2++c #t <
02 4 á
where o² = Var [l(e, fe( t))]]. Let a = 0
1-6
eq. (13)
gives
1- 6
Pr L L22
2
(14)
and therefore, with probability 1-5, the condition-speccficc loss is bounded by pto V V
Although
do not control the variance, I μ empirically (see appendix B.1 for details of standard error analysis) and therefore, the condition-specific loss is bounded by
Eg [e(c, (e( t)) 0
(15)
1 1+
8
(p))
therefore the per-condition loss is ((p) with high probability.
B ADDITIONAL EXPERIMENTS RESULTS
B. EMPIRICAL RESULTS AMORTIZATION LOSS AND ERROR ANALYSIS
To support the probabilistic generalization argument in appendices A.I and A.2 we plot the trajec- tory of the training amortized loss mean with its standard deviation over mini-batches throughout
15

---

Published as a conference paper at ICLR 2026
3.0
0.0
-0.-
2.0
-1.0-
1.0
-1.
-2.0
2000
4000
6000
8000
10000
12000
14000
16000
4000
6000
8000
10000 12600 14000 16000
Batches
Batches
Figure 4: (left) The training loss of each batch is consistently above the standard deviation. (right) The difference between the standard deviation and the mean confirms that the standard devi- ation is consistently less than equal to the mean.
Table 4: Ablation studies with 4 different architectural setups on the breast cancer dataset; test score the average over all proteins.
Method
Best Protein Predictor
0.489
Single channel
0.667
Multi channel
0.596
w/o output channel attention
0.581
w/o UNet channel attention
0.541
Condition injection w/ clement-wise addition
0.516
Condition injection w/ channel attention
0.535
Base (unconditional)
-0.017
training. The goal to empirically validate the bounded loss assumption (Assumption T and the assumption that the standard deviation of the loss remains small relative to the mean.
As shown Figure 4 the empirical standard deviation remains consistently lower than the mean across all training batches. This indicates that the loss distribution across conditioning configurations stable and tightly concentrated, justifying our use of probability bounds on the condition-specifi loss based amortized loss. This empirical observation confirms that the generalization bound derived in appendix AA2 holds in practice, with the per-condition loss reliably bounded within small multiple of the amortized training loss.
B.2 ABLATION STUDIES
We conduct ablation experiments disentangle the contributions of each architectural component. Using the breast cancer IMC dataset, we compare two feature injection mechanisms: elementwise addition and soft channel-wise attention, as well three attention configurations: (1) with channel- wise attention only in the UNet blocks, (2) with channel-wise attention only in the output block, and (3) the full model combining both mechanisms.
Table summarizes the contributions of each architectural component. The unconditional baseline performs poorly, with Pearson r < 0.1. Adding hierarchical feature injection significantly improves performance, and conditional injection via soft channel attention further improves over conventional element-wee wiieent-wiieeee addition. Output space channel attention provides modest gains, whereas latent- channel attention achieves more substantial improvements. capturing complex inter-channel depen- dencies in the latent feature space. These results support the integration of both spatially aligned conditioning and channel-wise attention for multi-channel imputation.
16

---

Published as a conference paper at ICLR 2026
B.3 ADDITIONAL RESULTS AND VISUALIZATION SPATIAL PROTEOMICS AND MRI GENERATION SAMPLES
We performed empirical uncertainty analysis. For each of 500 test cases, generated 5 in- dependent samples and computed pixel-wise variance as a measure of predictive uncertainty. The mean variance 0.003, indicating highly consistent predictions. In addition, we computed local Pearson correlation with a sliding 5x5 window to quantify spatial accuracy. Together, these form a pixel-level accuracy-uncertannyy map (Figure (5). where normalized variance ranges from and Pearson r from -1 to 1. The empirical distribution concentrates in the high-correlation, low- variance region (top right), showing that most pixels are both accurate and confident. Within this region, we observe the expected negative relationship: as local correlation decreases, variance sys- tematically increases, providing a well-calibrated link between predictive accuracy and uncertainty.
0.040
0.035
0.030
Tiled pearson
0.025
Percentage loading of normalized variance levels at each correlation level
0.020
0.015
0.010
0.005
0.000
Variance/Uncertainty level
Figure 5: Visualizations of the relationship of generation accuracy and sample uncertainty.
To assess whether our model captures biologically meaningful relationships between protein chan- nels, we evaluated the correlation between channels of the generated images. Specifically, we com- puted the pairwise Pearson correlation matrix across channels on (1) the generated outputs and (2) the ground truth images from the original data. We then compared these correlation matrices to assess how well the joint channel distribution is preserved. The Pearson correlations between the generated and ground-truth correlation matrices are 0.925 (MSE 0.031) in the lung dataset and = 0.892 (MSE 0.037) in the breast dataset. We found that the correlation structure in the generated data closely mirrors that of the ground truth, demonstrating that our model not only accu- rately imputes each channel but also retains the inter-channel dependencies critical for downstream biological interpretation.
Figure 6 shows high-resolution generated samples of the nucleus marker (DNA1), tumor markers (PanCK, Ki67, MMP9), immune markers (CD3, CD4, CD8a, CD45RO, FOXP3, CD68), and stro- mal markers (CD31, CD140b), confirming the generation quality and fidelity of our model.
Figure 7 shows the generated MRI samples for all 4 MRI signals (TICE, TIN, FLAIR, and T2W), suggesting our model generalizes well to other spatial biological data.
17

---

Published as a conference paper at ICLR 2026
Table 5: Comparison of predictive performance across ControlNet and BrushNet baselines, well
as hybridized variants in which control modules are reused with our denoising backbone.
Method
Breast
Lung
ControlNet (Zhang et al.
2023
0.452
0.537
BrushNet (Ju et al. 2024
0.357
0.394
ControlNet hybrid
0.516
0.583
BrushNet hybrid
0.486
0.541
Ours (single channel) )
0.667
0.703
Ours (multi channel)
0.596
0.647
B.4 HYBRIDIZATION EXPERIMENT OF BENCHMARKING DIFFUSION-BASED ALGORITIMMS
To better compare the performance of our model with the SOTA diffusson-based method, we test both ControlNet and BrushNet as the baselines, with additional experiments with hybrid architec- tures, where we combine the conditional network in ControlNet (Zhang et al. 2023) and BrushNet (Ju al. 2024 with our diffusion branch.
We evaluate each model on two IMC datasets (Breast and Lung), measuring the performance with Pearson r correlation. "ControlNet hybrid" refers to model that reuses the ControlNet condition- ing branch but replaces the main denoising network with our backbone, which includes the channel attention mechanism we proposed. Similarly, "BrushNet hybrid" would replace BrushNet's denois- ing path with ours, in which the model has spatial attention in the UNetBlocks. Both hybrid models use element-wise addition for injection.
As shown in Table 5 ControlNet outperforms BrushNet. Notably, BrushNet's UNetBlock does not have either a channel or spatial attention module, whereas ControlNct incorporates spatial attention. The extra spatial attention may help the performance of ControlNet, even though BrushNet claims to have a better performance natural image tasks. However, BrushNet's performance is not ideal in the hybrid setup, which introduces spatial attention in the UNetBlocks, suggesting over- conditioning. Our full model, trained jointly using random masking, outperforms both baselines and hybrid models by a large margin.
C MODEL ARCHITECTURE AND EXPERIMENTS SETUP
C.I MODEL ARCHITECTURE
We implement model based the EDM framework (Karras et al. 2022), modified and adapted for conditional generation under the multichannel setup. The diffusion branch has a 3-level UNet with downsampling and upsampling paths. Each UNetBlock contains a spatial attention module followed by a channel attention module. Time conditioning is implemented via sinusoidal positional embeddings.
Spatial conditions are processed by the conditional network, which has an architecture identical to the diffusion branch. Conditional features are injected into the diffusion network in middle blocks and skip connections, following ControlNet, but with soft channel attention instead of element- wise addition. We also implement the model such that it can inject the conditional features each resolution level, like BrushNet, but as shown in table the performance is suboptimal. The output layer contains a followed by the output channel attention layer.
convolutional network to map from the latent space to the protein space,
C.2 DATA PROCESSING
The raw multichannel IMC data are processed using the Yeo-Johnson transformation adjust the signal strength and normalized to j0. 1] by the 1st and 99th percentiles of each channel. The processed IMC data are partitioned into 64 x 64 patches. Note that our model only works data in the original data space since we need the alignment between the same channel across multiple datasets for the multi-dataset joint training. With pre-trained MedVAE (Varma al. 2025) finetuned IMC
18

---

Published as a conference paper at ICLR 2026
single-channel images, which encode single-channel images to lower resolution, can compress the size of the IMC images by a factor of 4 for each channel independently. By doing so, we get C' x 16 16 IMC data where each channel maintains its original protein semantics. The training and test splits are grouped by different patients.
The single-cell CITE-seq data are also processed with the Yeo-Johnson transformation to adjust the signal strength and normalized to [0, 1] by the 1st and 99th percentiles of each channel.
C.3 LOSS AND TRAINING SETUP
The model trained based on the pise-prediction objective. For each sample Do. we generate noisy version q( (oo) and learn to predict the corresponding noise c, given the spatial condition and timestep The training loss is:
Pr. [lle cs (z,, t, c)1[3]
(16)
where C denotes random subset of observed channels.
For single-channel generation tasks, one target channel is masked out, and loss is evaluated on the target channel only. For multichannel generation, the channels are sampled by independent Bern(0.9) and the loss is evaluated on the generated full-channel data. The general training hyper- parameters are:
Optimizer: Adam with , = (0.9, 0.999)
Learning rate: x 10-1
Batch Size: 256
Noise Schedule: EDM noise scheduling
D COMPUTATION RESOURCES
All experiments were trained on NVIDIA A5000 GPUs with 24G RAM. The model was trained for 2000k images with batch size of 256, taking approximately 2 hours complete with 16 x 16 resolution. All results were obtained using a single-GPU setup unless otherwise specified.
All models were implemented in PyTorch based on the EDM framework (Karras et al. 2022 and trained using standard FP32 precision without mixed-precision.
E DISCLOSURE AND STATEMENTS
LARGE LANGUAGE MODEL USAGE
LLMs used to check grammar and spelling and polish wording.
E.2 ETHICS STATEMENT
All authors of submitted papers have read the Code of Ethics and adhere to it.
19

---

Published as a conference paper at ICLR 2026
Ground Truth
Generated
Ground Truth
Generated
DNA1
PanCK
Ki67
FOXP3
CD3
CD45RO
CD4
CD8a
CD31
CD140b
CD68
MMP9
Figure 6: Visualizations of generated spatial proteomics samples of tumor and immune marker proteins of the tumor microenvironment.
20

---

Published as a conference paper at ICLR 2026
TICE
FLAIR
T2W
predicted TICE
predicted T1N
predicted FLAIR
predicted
heldout channel
Figure 7: Visualizations of generated MRI samples for each of the MRI signals.
21


Reviews:
[{"review_id": "KhQmNvvvYH", "reviewer": "Reviewer_WNmC", "summary": "The authors introduce a Multi-Channel Diffusion (MCD) framework, designed for controllable synthesis and imputation of multi-channel biological data such as spatial proteomics, single-cell omics, and MRI. This framework integrates a mechanism for hierarchical spatial feature injection with a dual channel-attention module. This enables the resulting models to preserve spatial alignment while capturing complex inter-channel dependencies. A wide range of experiments on publicly available benchmarks suggest that MCD outperforms existing baselines.", "questions": "I would ask the authors to address the weaknesses highlighted above. Particularly:\n1) Evaluate differences in predictive performance between their method and the others through appropriate statistical tests\n2) Check whether expected cell-specific markers are expressed in the inputed data. You could also scale up this assessement to the pathway level, and check which biological pathways are enriched in the imputed data and if the enrichment results are biologically meaningful", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1) This paper combines two different methodological innovations in a quite ingenious way, effectively addressing the spatial and inter-channel complexity of biological data.\n\n2) The resulting models demonstrates versatility across multiple domains, including spatial proteomics, single-cell omics, and MRI modality synthesis, showing strong generalization and scalability.\n\n3) Finally, the presented evaluation is quite comprehensive. I particularly appreciate the ablation studies that assess the individual contributions of model components.", "weaknesses": "1) All comparisons reported in Tables 1 to 3 lack any assessement of stastistical significance. This makes it difficult to gauge whether differences in performances are actually significant.\n2) There is not biologically-grounded evaluation of the imputed data. For example, are known protein markers expressed in their corresponding cells?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hdi5pEbz7k", "reviewer": "Reviewer_aLRt", "summary": "This work proposes a unified multi-channel diffusion (MCD) framework for controllable generation of structured biological data.", "questions": "1. The authors declared the following: \"All experiments were trained on NVIDIA A5000 GPUs with 24 GB of VRAM. The model was trained for 2000k imgs with a batch size of 256, taking approximately 2 hours to complete at 16 × 16 resolution. All results were obtained using a single-GPU setup unless otherwise specified.\" Which kind of biological data have a resolution of 16 × 16?", "limitations": "", "rating": null, "confidence": 4, "soundness": 1, "presentation": 1, "contribution": 1, "strengths": "1. The idea of developing a framework capable of controllably generating multi-channel biological data using diffusion models is interesting.", "weaknesses": "1. The paper is quite obscure and its objective remains unclear. The title suggests that it focuses on developing a generative framework for multi-channel biological data, but the type of data is not specified. I assumed the authors were referring to images, yet in the experiments they attempt to predict protein expression from paired scRNA-seq data, and later they evaluate their method on MRI images. This inconsistency makes the overall methodology difficult to understand and significantly undermines the coherence of the paper.\n\n2. The paper contains several incorrect claims and assertions. To list a few: “Existing generative models typically assume low-dimensional inputs (e.g., RGB images)” and “In spatial profiling data, each channel designates a specific molecule of interest (e.g., proteins n ≥ 30 and genes n ≥ 100), and each pixel (or cell)…”.\n\n3. The manuscript contains numerous writing inconsistencies, redundant and buzzword-heavy claims, and incorrect or oversimplified descriptions of diffusion theory. Core method elements (random channel masking, SE attention) are incremental and poorly justified as novel contributions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JGl0dLFmd3", "reviewer": "Reviewer_iF6S", "summary": "The paper proposes MCD, a conditional diffusion framework for multi-channel biological data. Conditioning is handled via masked images. The proposed architecture is a dual network with hierarchical feature injection from a contextual encoder into the denoiser, plus two channel-attention mechanisms (SE for injection and transformer-style channel self-attention inside unet blocks). Experiments show state of the art level results on CITE-seq protein prediction, IMC channel imputation including a union vs intersection multi-dataset study, hybrid controls versus ControlNet/BrushNet, and BraTS MRI.", "questions": "See the questions and actionable items below. I find the core idea is strong, but the draft feels rushed; clearer presentation and a few added analyses would better show the paper's potential.\n\n\n- **MedVAE ablation.** It would be nice to see the effect of the latent space choice on the results. One can convert each channel into grayscale (or stack each channel 3 times to imitate RGB) and run through a stable diffusion VAE pipeline do the latent diffusion.\n\n- **Efficiency plots.** Single step generation has a clear advantage from its strong results and efficiency but it would be nice to see more details on how you distilled together with the performance comparisons with other baselines.\n\n- **Confidence intervals.** Report mean $\\pm95\\%$ CI over multiple seeds for all main tables to quantify variance from training and random masking. \n\n### Mistakes in text.\n\n- **Section 2.2** The part describing classical RGB image imputation and classical RGB colorization problems are not correct. If we have $C_m = 0$ and $C_o = C = 3$ there is nothing to impute. Additionally, for the colorization example $C_m = 3$ and $C_o = 1$ implies $C = C_m + C_o = 4$ which is clearly not RGB. \n\n- **Dataset reference for BMMC.** Table 1 report BMMC (bone marrow mononuclear cells) as BMNC both in the body and in the caption.\n\n- **Different acronym on Table 3.** Throughout the paper the authors refer their method as MCD however Table 3 uses DiffuseMRI", "limitations": "", "rating": 2, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 3, "strengths": "- **Problem relevance.** Training with random channel masking yields one model that accepts arbitrary observed subsets and making it flexible. The union and intersection result supports cross-dataset integration under partial channel overlap.\n\n- **Strong empirical results.** When reported, the method consistently outperforms baselines. Experiments are broad and span single/multi dataset setups and including hybrid controls.\n\n- **Ablations.** Stepwise ablations and ControlNet/BrushNet hybrids help attribute gains to hierarchical injection.", "weaknesses": "- **Subset-size stress-tests are missing.** One of the core claims is robustness to arbitrary observed subsets, but there is no sweep of performance vs. #observed channels / masking-probability p, nor targeted leave a group out per channel families. Single vs multi channel and union vs intersection is positive but partial.\n\n- **Efficiency evidence.** Table 1 lists SiD(1-step) with near identical accuracy and claims two orders of magnitude speedup, but there are no wallclock analysis for readers to observe the real performance gains. \n\n- **Method clarity.** I find the pieces in the methods section hard to put together. It requires stitching together several sections to reconstruct the exact flow of input, context encoder, SE-gated injections per scale, denoiser block attention, output SE flow. A single forward schematic/pseudocode can improve the clarity significantly. \n\n- **CFG analogy is conceptual.** I could not directly link random masking with CFG. Instead, the authors can provide a deeper the analysis of how the compared baselines actually operate and where they fall short.\n\n- **Reproducibility.** The paper promises code upon publication without anonymous repo or supplementary materials. This limits the verification significantly.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aoilMLBXfc", "reviewer": "Reviewer_t6Ur", "summary": "Many biological modalities are inherently multi‑channel and spatially co‑registered (IMC, clinical imaging, spatial transcriptomics). Off‑the‑shelf conditional diffusion often assumes low‑dimensional channels and crude conditioning that breaks spatial correspondence. The paper proposes MCD (multi‑channel diffusion) for controllable generation and imputation given arbitrary subsets of observed channels. Core claims: (1) hierarchical feature injection to keep spatial alignment; (2) two channel‑attention modules to model inter‑channel dependencies; and (3) random channel masking for amortized training across any condition configuration. Experiments cover single‑cell gene‑to‑protein prediction, IMC imputation, and cross‑dataset generalization.", "questions": "- (Multi-channel vs multi-source) The setting is similar to multi-source integration problem, and the masking idea is similar to scVAEIT. It would be interesting to discuss the connection. (Du, J.‑H., Cai, Z., Roeder, K. (2022), Robust probabilistic modeling for single‑cell multimodal mosaic integration and imputation via scVAEIT.”)\n- (Masking policy) What masking probability $p$ and schedule work best? Any performance cliffs when masks are extremely sparse/dense? Show sensitivity curves and OOD mask combos not seen in training.\n- (Scalability) Provide computational complexity (e.g., wall‑clock, memory, and FLOPs) versus $C$, $D$, and $H\\times W$ for both attention modules. Where does channel self‑attention become the bottleneck?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The hierarchical feature‑injection pipeline preserves spatial correspondence at every resolution, which is exactly what biological channels need. The injection is simple to deploy and avoids brittle concatenation tricks common in vision diffusion.\n- Pairing SE gating with transformer‑style channel attention gives a plausible division of labor: per‑sample recalibration and higher‑order inter‑channel structure. The additional SE head at the output enforces cross‑channel coupling where it matters.", "weaknesses": "- Comparators underpowered for spatial tasks. ControlNet is a fair reference, but the paper doesn’t benchmark strong end‑to‑end, jointly‑trained spatial conditioners or ablate cross‑attention conditioning versus the proposed SE‑gated injection. Given that the method already computes $E_{\\ell}(c)$, a direct cross‑attention baseline would be informative.\n- Masking policy and distribution shift. Random masking zeros unobserved channels in the condition (Algorithm 1). The paper does not analyze whether zero‑masking creates a train/test mismatch when “missing” at test time means “physically unmeasured but nonzero distribution,” especially for modalities where absolute intensity carries semantics.\n- Metrics don’t reflect calibration. Spatial evaluation is Pearson r. No per‑channel calibration (e.g., bias/variance decomposition), no uncertainty quality, no region‑level histograms for key biomarkers. For clinical‑adjacent imaging, correlation alone can be misleading.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)