Joint Model-based Model-free Diffusion (JM2D)

• The paper introduces JM2D as a design principle that couples explicit model-based structure with model-free generative diffusion to reduce hallucinations and enforce physical or operational constraints.
• In mixed-weather image restoration, JM2D leverages atmospheric physics and binary degradation masks to guide conditional denoising, yielding improved PSNR and SSIM performance.
• For planning with constraints, JM2D employs joint sampling of diffusion-generated action sequences and model-based optimization outputs to ensure feasibility and reduce intervention rates.

Joint Model-based Model-free Diffusion (JM2D) denotes a class of diffusion-based methods that couple explicit model-based structure with model-free generative denoising. In the literature considered here, the designation covers two technically distinct but conceptually aligned settings: mixed-weather image restoration, where a Joint Conditional Diffusion Model (JCDM) combines an atmospheric-physics degradation model, degradation masks, conditional diffusion, and uncertainty-aware refinement; and planning with constraints, where a pre-trained diffusion planner is combined with a model-based optimization module through joint sampling under an interaction potential (Yue et al., 2024, Jung et al., 10 Sep 2025). In both settings, the central aim is to preserve the expressive prior of diffusion while constraining generation by physically or operationally meaningful structure.

1. Terminology and scope

In mixed-weather image restoration, the problem is blind restoration under mixed weather degradations, where a single image may simultaneously contain rain streaks, adherent raindrops, haze or veiling, and snow. The supplied account identifies three difficulties: different degradations overlap spatially and interact nonlinearly; accurate per-component estimation is brittle and errors cascade to the final reconstruction; and severe corruption removes texture and fine details, leaving an ill-posed inverse problem. Diffusion alone is described as prone to generating plausible but irrelevant content when observation guidance is weak, whereas physical constraints narrow the solution space and spatial degradation masks focus denoising capacity on corrupted regions (Yue et al., 2024).

In planning with constraints, the problem is the integration of a learned, model-free diffusion planner with a model-based optimization module that enforces mission- and safety-critical constraints. The supplied account states that naive post-hoc integration can be incompatible because diffusion’s multimodal samples may be adversarial to an optimization module with limited authority or feasibility, producing overrides, infeasible corrections, or degraded task performance. JM2D addresses this by formulating integration as a joint sampling problem rather than a sequential correction pipeline (Jung et al., 10 Sep 2025).

A plausible implication is that JM2D is best understood not as a single architecture, but as a design principle: diffusion is retained as the source of rich priors and multimodality, while model-based structure is introduced as a compatibility mechanism that reduces hallucination, infeasibility, or other forms of unconstrained generation.

2. Common formulation and the meaning of “joint”

Across the two instantiations, “joint” refers to coupling the diffusion process with external structure at the level of the generative variables themselves rather than treating the model-based component as a purely post-hoc filter (Yue et al., 2024, Jung et al., 10 Sep 2025).

Setting	Model-based component	Model-free component
Mixed-weather restoration	Atmospheric scattering model, atmospheric light $A$, transmission $t(x)$, binary masks $M(x)$	Conditional diffusion model and refinement U-Net with UEB
Planning with constraints	Optimization module over $k$, hard constraints $g(\tau)\le 0$, equalities $h(\tau)=0$, interaction potential $V(x,k)$	Pre-trained diffusion prior over action sequences $x$

In restoration, jointness is realized through argument-level conditioning: the degraded image $y=I(x)$ and the predicted binary mask $m(x)$ are both fed into the denoiser as $t(x)$0. The degraded observation anchors the generative process to observed content, while the mask focuses denoising on corrupted regions. The supplied description explicitly notes that this is compatible with concatenation-based or feature-modulation implementations, although the instantiated design uses direct inputs to $t(x)$1 (Yue et al., 2024).

In planning, jointness is formalized probabilistically. Let $t(x)$2 denote the action sequence produced by the model-free planner and $t(x)$3 the model-based module’s output. JM2D defines a joint target

$t(x)$4

where $t(x)$5 is the pre-trained diffusion prior, $t(x)$6 is a model-based prior, and $t(x)$7 is an interaction potential that scores mutual compatibility. A common choice is

$t(x)$8

The same source states that conditional generation is a special case of JM2D when the interaction factorizes and the $t(x)$9-prior is uniform (Jung et al., 10 Sep 2025).

The common structural motif is therefore not merely “physics plus diffusion” or “optimization plus diffusion,” but a tighter coupling in which the model-based side shapes the conditional or joint distribution sampled by diffusion.

3. Mixed-weather image restoration as JM2D

The image-restoration instantiation is centered on the paper “Joint Conditional Diffusion Model for Image Restoration with Mixed Degradations” (Yue et al., 2024). Its starting point is a mixed degradation model derived from classic single-degradation formulations for rain streaks, adherent raindrops, snow, and haze. The haze term follows the atmospheric scattering model (ASM),

$M(x)$0

where $M(x)$1 is the clean image, $M(x)$2 the observed image, $M(x)$3 the global atmospheric light, $M(x)$4 the scattering coefficient, and $M(x)$5 the depth.

Motivated by ASM and the rain veiling effect, the supplied description recasts rain streaks and raindrops using atmospheric light and binary masks. This aligns rain streaks, raindrops, and snow with an $M(x)$6-blending form analogous to haze. The resulting mixed-degradation synthesis is written as

$M(x)$7

with $M(x)$8 and

$M(x)$9

Here, $k$0 collects binary masks for rain streaks, raindrops, and snow. For training synthesis, rain-streak and snow masks are drawn from Rain100H and Snow100K, while raindrop masks are generated by a metaball model to vary shape, size, and location. At inference, a mask prediction branch provides a binary degradation mask $k$1 from the input image.

The model-free side is a conditional denoising diffusion process. Given clean $k$2, the forward noising chain is

$k$3

with

$k$4

The learned reverse process is conditioned on the degraded image $k$5 and predicted mask $k$6:

$k$7

and training uses the standard $k$8-prediction objective

$k$9

Sampling is accelerated by an implicit deterministic sampler described as DDIM-style, with $g(\tau)\le 0$0 deterministic steps at inference rather than the full $g(\tau)\le 0$1 noising steps used in training. After diffusion yields a coarse restoration, a refinement U-shaped network reconstructs the restored image and incorporates an Uncertainty Estimation Block (UEB) at each scale. The UEB models aleatoric uncertainty $g(\tau)\le 0$2 and epistemic uncertainty $g(\tau)\le 0$3, approximating pixel-wise uncertainty as $g(\tau)\le 0$4. Feature modulation at scale $g(\tau)\le 0$5 is

$g(\tau)\le 0$6

and the decoder outputs the final refined restoration $g(\tau)\le 0$7.

Training uses Adam with $g(\tau)\le 0$8, $g(\tau)\le 0$9, batch size $h(\tau)=0$0, input patches $h(\tau)=0$1, and learning rate $h(\tau)=0$2. The refinement stage is trained with

$h(\tau)=0$3

an uncertainty-aware term

$h(\tau)=0$4

and the total loss

$h(\tau)=0$5

Within the supplied description, this construction is explicitly identified as a JM2D method because the atmospheric-physics mixed-degradation model governs synthesis and guidance, while the conditional diffusion prior and refinement network perform restoration beyond purely deterministic inversion.

4. Planning with constraints as JM2D

The planning instantiation appears in “Joint Model-based Model-free Diffusion for Planning with Constraints” (Jung et al., 10 Sep 2025). The setup uses state $h(\tau)=0$6, action $h(\tau)=0$7, a finite-horizon trajectory $h(\tau)=0$8, and environment dynamics $h(\tau)=0$9. The planner produces an action sequence $V(x,k)$0, while the model-based module outputs $V(x,k)$1. Hard safety and feasibility constraints are encoded as $V(x,k)$2 and equalities $V(x,k)$3.

Given a diffusion plan $V(x,k)$4, the model-based module solves

$V(x,k)$5

The critique of standard compositions is twofold. First, post-hoc filtering samples $V(x,k)$6 and only then corrects via $V(x,k)$7, which can fail when the sampled plan lies far from the optimizer’s feasible region. Second, gradient-guided and projection-based methods operate on noisy states, can be brittle for non-differentiable or non-convex modules, and may harm data fidelity.

JM2D instead defines a joint diffusion process over $V(x,k)$8. The forward noising process is

$V(x,k)$9

with independent schedules $x$0. Reverse-time denoising requires the joint score $x$1, which is intractable because $x$2 acts on clean variables and may be non-differentiable. The supplied derivation uses a Tweedie-form identity and approximates the score by self-normalized importance sampling.

With proposal

$x$3

and

$x$4

the practical estimator is

$x$5

The account emphasizes that this requires only function evaluations of $x$6 on clean samples, without gradients of the optimizer and without retraining.

The high-level algorithm initializes $x$7, constructs Monte Carlo clean samples at each diffusion level, computes importance weights

$x$8

estimates the joint score, and applies a DDIM-style update. Optional resampling can monitor effective sample size. If the final draw fails the hard constraint, a single model-based correction is applied. The paper further states that conditional generation is recovered as a limiting case when the interaction factorizes appropriately.

5. Empirical behavior and computational profile

The two JM2D instantiations are evaluated on distinct benchmark families, but both are reported to improve alignment between diffusion priors and model-based structure (Yue et al., 2024, Jung et al., 10 Sep 2025).

For mixed-weather restoration, JCDM with refinement is reported across six Cityscapes-based synthetic joint-degradation cases: case 1 (rain streak) achieves $x$9 dB and $y=I(x)$0 SSIM; case 2 (rain + snow) $y=I(x)$1; case 3 (rain + light haze) $y=I(x)$2; case 4 (rain + heavy haze) $y=I(x)$3; case 5 (rain + moderate haze + raindrop) $y=I(x)$4; and case 6 (rain + snow + moderate haze + raindrop) $y=I(x)$5. The supplied comparison states that JCDM surpasses WeatherDiff and is competitive or superior to IRNeXt across mixed cases. An ablation attributes a substantial part of the gain to joint conditioning on the mask $y=I(x)$6: case 4 improves by $y=I(x)$7 dB PSNR and $y=I(x)$8 SSIM, case 5 by $y=I(x)$9 dB and $m(x)$0, and case 6 by $m(x)$1 dB and $m(x)$2 over conditioning on $m(x)$3 alone.

For single-weather restoration, the reported results are as follows.

Benchmark	JM2D / JCDM result	Comparison stated
Raindrop dataset	$m(x)$4 dB / $m(x)$5 SSIM	Surpasses All-in-one $m(x)$6 and WeatherDiff $m(x)$7; AIRFormer has higher SSIM $m(x)$8 but lower PSNR $m(x)$9
Dense-Haze	$t(x)$00 dB / $t(x)$01 SSIM	Comparable or better than baselines such as FocalNet $t(x)$02
NH-HAZE	$t(x)$03 dB / $t(x)$04 SSIM	PSNR slightly higher than AIRFormer $t(x)$05; FocalNet has higher SSIM $t(x)$06
Snow100K-L	$t(x)$07 dB / $t(x)$08 SSIM	Best PSNR and tied best SSIM with IRNeXt $t(x)$09

The computational profile for $t(x)$10 images is also explicit. JCDM without refinement uses $t(x)$11M parameters, $t(x)$12G FLOPs, and $t(x)$13 s latency; with refinement it uses $t(x)$14M parameters, $t(x)$15G FLOPs, and $t(x)$16 s latency. The supplied description states that this is over $t(x)$17 faster inference than WeatherDiff at $t(x)$18 s while achieving superior quality.

For planning, the offline RL evaluation uses D4RL PointMaze with a reachability-based safety filter. Under wall padding $t(x)$19, the selected table reports, for a high-quality backup planner, RAIL at Safe Success $t(x)$20, Task Horizon $t(x)$21, and Intervention Rate $t(x)$22, whereas JM2D achieves $t(x)$23, $t(x)$24, and $t(x)$25. Under degraded backup planners, RAIL’s Safe Success drops to $t(x)$26 and $t(x)$27, while JM2D remains at $t(x)$28 and $t(x)$29, with lower intervention. The text also states that no method violates safety.

In the conditional-generation “Donut” domain, JM2D produces feasible, in-distribution samples with $t(x)$30 valid, while projection-based conditional diffusion yields $t(x)$31 samples far from data. In D3IL Avoiding, the selected table reports that JM2D attains Safe Success $t(x)$32 with $t(x)$33 violations in the Top-Right scenario, $t(x)$34 with $t(x)$35 violations in the Both scenario, and $t(x)$36 with $t(x)$37 violations in the Cluttered scenario, outperforming SafeDiffuser and MPD in those settings. An ablation on clean-sample estimation shows that increasing partial denoising depth $t(x)$38 from $t(x)$39 to $t(x)$40 improves Safe Success and reduces violations. On a real Franka mug-pickup task with unseen obstacles, the supplied description states that vanilla diffusion collides, whereas JM2D with a reachability safety filter completes the task safely.

The computational trade-off in planning is different from restoration. The method increases cost versus vanilla diffusion because Monte Carlo clean-sample construction is performed at every step. The stated practical mitigations include partial denoising with small $t(x)$41, candidate subsampling, and faster ODE solvers such as DPM-Solver++.

6. Limitations, misconceptions, and future directions

The supplied sources identify distinct failure modes in the two JM2D settings (Yue et al., 2024, Jung et al., 10 Sep 2025). In image restoration, performance is sensitive to mask quality: masks are effective guidance, but errors in the predicted degradation mask can degrade restoration quality. The physics model is rooted in ASM and binary masks, whereas real-world weather may involve spatially varying atmospheric light, specularities, or dynamic droplets; this potentially limits exact physical fidelity. The reported results also show metric trade-offs, since some competing models achieve higher SSIM on certain datasets even when JCDM attains higher PSNR. The stated future direction is to refine the diffusion process to better preserve semantic details while removing severe degradations.

In planning, the method relies on support overlap: the feasible subset must have non-zero mass under $t(x)$42, otherwise importance sampling degeneracy occurs. Hard constraints are not guaranteed for every draw; the procedure steers sampling toward feasibility, and a single model-based post-check or correction is used when strict guarantees are required. Weight collapse can occur if the interaction potential is too peaky or $t(x)$43 is too large, and computational cost scales with the diffusion steps, clean-sample depth, and candidate counts.

Several common misconceptions are directly countered by the supplied material. JM2D is not merely sequential post-hoc filtering: in planning it is posed as a joint sampling problem, and in restoration the denoiser is jointly conditioned on the degraded image and degradation mask. It is also not restricted to differentiable constraints: the planning formulation is explicitly designed for non-differentiable and non-convex modules because it needs only evaluations of $t(x)$44 on clean samples. Conversely, JM2D is not a universal guarantee of correctness; both instantiations remain dependent on the quality of their model-based side, whether that is mask prediction and atmospheric assumptions in restoration or support overlap and importance-sampling stability in planning.

Taken together, the literature suggests that JM2D is a general strategy for constraining diffusion by structured external knowledge without discarding multimodality. In restoration, that structure is atmospheric physics and spatial masks; in planning, it is compatibility with a constrained optimization module. The common contribution is to shift the role of the model-based component from after-the-fact correction to an integral part of the generative process.