Flexi-UNet: Scale-Adaptive Diffusion

• Flexi-UNet is a scale-adaptive UNet variant that integrates scale-space theory into DDPM pipelines, enabling dynamic routing across variable resolutions.
• It optimizes computational efficiency by selectively activating encoder and decoder blocks, reducing FLOPs and training time while maintaining generation quality.
• Empirical evaluations and ablation studies demonstrate Flexi-UNet’s ability to balance perceptual fidelity and resource efficiency through explicit resolution control.

Flexi-UNet is a UNet variant introduced to efficiently support the Scale Space Diffusion (SSD) framework, enabling dynamic routing and computation at variable resolutions throughout the generative diffusion process. It modifies the canonical denoising diffusion probabilistic model (DDPM) pipelines by integrating scale-space theory, which offers explicit control over resolution schedules and network path traversal, thereby optimizing computational resources during both training and inference without substantial degradation in generation quality (Mukhopadhyay et al., 9 Mar 2026).

1. Mathematical Foundations in Scale Space Diffusion

Flexi-UNet is implemented in the context of the SSD framework, a generalization of DDPM achieved by replacing the traditional scalar noise schedule’s multiplier with a family of linear operators $M_t$. This allows incorporation of degradations such as downsampling or blurring; specifically,

$$M_t = \frac{a_t}{a_{t-1}}\;\mathrm{Downsample}_{r(t)/r(t-1)}$$

where $r(t)$ is a monotonically decreasing resolution schedule, and $x_0$ is sampled from the data distribution at full resolution. The noising process for each timestep is governed by: $$q(x_t \mid x_{t-1}) = \mathcal{N}(x_t \,;\, M_t\,x_{t-1},\,\Sigma_{t\mid t-1})$$ with covariance

$$\Sigma_{t\mid t-1} = \sigma_t^2I - \sigma_{t-1}^2\,M_t\,M_t^\top$$

and the closed-form marginal: $$x_t = M_{1:t}\,x_0 + \sigma_t\,\epsilon, \quad \epsilon \sim \mathcal{N}(0, I)$$ The reverse (denoising) step leverages the tractable Gaussian posterior

$$q(x_{t-1} \mid x_t, x_0) = \mathcal{N}(\mu_{t\to t-1}, \Sigma_{t\to t-1})$$

where

$$\Sigma_{t\to t-1} = \sigma_{t-1}^2I - \frac{\sigma_{t-1}^4}{\sigma_t^2}M_t^\top M_t$$

$$\mu_{t\to t-1} = M_{1:t-1}\,x_0 + \frac{\sigma_{t-1}^2}{\sigma_t^2}M_t^\top(x_t - M_t\,M_{1:t-1}\,x_0)$$

In practice, $$M_t = \frac{a_t}{a_{t-1}}\;\mathrm{Downsample}_{r(t)/r(t-1)}$$0 is replaced by the Flexi-UNet model’s prediction at the appropriate lower resolution. Training proceeds with a Min-SNR-$$M_t = \frac{a_t}{a_{t-1}}\;\mathrm{Downsample}_{r(t)/r(t-1)}$$1 weighted loss: $$M_t = \frac{a_t}{a_{t-1}}\;\mathrm{Downsample}_{r(t)/r(t-1)}$$2 where $$M_t = \frac{a_t}{a_{t-1}}\;\mathrm{Downsample}_{r(t)/r(t-1)}$$3 is the standard SNR.

2. Flexi-UNet Topology and Dynamic Routing

Flexi-UNet modifies the standard ADM UNet to efficiently process variable input resolutions by dynamically activating only the relevant subset of encoder and decoder blocks. The maximum network depth is $$M_t = \frac{a_t}{a_{t-1}}\;\mathrm{Downsample}_{r(t)/r(t-1)}$$4; resolutions are indexed from coarsest ($$M_t = \frac{a_t}{a_{t-1}}\;\mathrm{Downsample}_{r(t)/r(t-1)}$$5) to finest ($$M_t = \frac{a_t}{a_{t-1}}\;\mathrm{Downsample}_{r(t)/r(t-1)}$$6), each corresponding to input resolution $$M_t = \frac{a_t}{a_{t-1}}\;\mathrm{Downsample}_{r(t)/r(t-1)}$$7. The dynamic routing is as follows:

• For input at $$M_t = \frac{a_t}{a_{t-1}}\;\mathrm{Downsample}_{r(t)/r(t-1)}$$8, the image is fed to encoder blocks $$M_t = \frac{a_t}{a_{t-1}}\;\mathrm{Downsample}_{r(t)/r(t-1)}$$9.
• After the bottleneck, decoder blocks $r(t)$0 are used.
• For a resolution-preserving step ($r(t)$1), the network stops at decoder level $r(t)$2 and outputs at resolution $r(t)$3.
• For a resolution-increasing step ($r(t)$4), an extra decoder is applied to reach $r(t)$5.

To address channel width mismatches, Flexi-UNet inserts 1×1 convolutions at every entry and exit point—these project the raw input into the feature dimensions expected by the encoder, and similarly, project features back to RGB at output. When a decoder block expects a skip connection from a bypassed encoder layer, a zero-tensor of the correct shape is used in place. Residual and attention weights, as well as time-embeddings, are shared identically with the baseline UNet.

Parameter scaling is efficient: Flexi-UNet adds only $r(t)$6 small 1×1 convolution layers, negligible relative to the main architecture. At maximum (full resolution), the network matches baseline UNet in FLOPs and parameters; at coarsest resolutions, it restricts computation to the innermost blocks, reducing both FLOPs and peak memory by roughly a factor of $r(t)$7.

3. Pseudocode for Forward and Reverse Processes

Forward (Noising) Pass

For $r(t)$8 to $r(t)$9:

1. Given input $x_0$0 at resolution $x_0$1, compute:

$x_0$2

Conditional Reverse (Denoising) Pass

1. Initialize $x_0$3 at coarsest resolution $x_0$4.
2. For $x_0$5 down to 1:
   • Use FlexiUNet$x_0$6 to predict $x_0$7 at $x_0$8.
   • Form posterior mean:

   $x_0$9

- Sample posterior noise: use Lanczos decomposition for non-isotropic $$q(x_t \mid x_{t-1}) = \mathcal{N}(x_t \,;\, M_t\,x_{t-1},\,\Sigma_{t\mid t-1})$$0 or standard Gaussian otherwise. - Set $$q(x_t \mid x_{t-1}) = \mathcal{N}(x_t \,;\, M_t\,x_{t-1},\,\Sigma_{t\mid t-1})$$1 posterior noise.

4. Empirical Evaluation

CelebA and ImageNet Results

Flexi-UNet was evaluated on CelebA at resolutions $$q(x_t \mid x_{t-1}) = \mathcal{N}(x_t \,;\, M_t\,x_{t-1},\,\Sigma_{t\mid t-1})$$2, $$q(x_t \mid x_{t-1}) = \mathcal{N}(x_t \,;\, M_t\,x_{t-1},\,\Sigma_{t\mid t-1})$$3, and $$q(x_t \mid x_{t-1}) = \mathcal{N}(x_t \,;\, M_t\,x_{t-1},\,\Sigma_{t\mid t-1})$$4, and on ImageNet at $$q(x_t \mid x_{t-1}) = \mathcal{N}(x_t \,;\, M_t\,x_{t-1},\,\Sigma_{t\mid t-1})$$5. Key metrics include FID, wall-clock training time, and per-iteration GFLOPs. Comparative results between DDPM-ε and Flexi-UNet at select resolution schedules are as follows:

Dataset@Resolution	Method	FID	Time (h)	GFLOPs/iter
CelebA@64×64 (2 levels)	Flexi-UNet (2 levels)	2.14	62.6	50.6
CelebA@64×64	DDPM-ε	2.22	70	60.05
CelebA@128×128 (3 lvls)	Flexi-UNet (3 levels)	6.53	31.7	87.4
CelebA@128×128	DDPM-ε	4.16	50.5	132.3
CelebA@256×256 (3 lvls)	Flexi-UNet (3 levels)	7.79	59.0	317.4
CelebA@256×256	DDPM-ε	5.52	87.3	497.0
ImageNet@64×64 (2 lvls)	Flexi-UNet (2 levels)	13.08	—	—
ImageNet@64×64	DDPM-ε	12.82	—	—

• At 2-level configurations, Flexi-UNet slightly improves FID, while reducing training time by approximately 10–20% and computational cost (GFLOPs) by 15–40%, depending on the number of levels.
• Increasing the number of levels reduces GFLOPs and time but increases the FID.
• Ablations show similar or improved inference time, especially with increased level granularity.

Schedules that dwell on higher resolutions yield improved FID but correspondingly increase training time—a tradeoff controlled by the resolution schedule (e.g., ConvexDecay_0.5 vs. Equal).

5. Architecture Ablations and Resolution Schedule Impact

Ablation studies compare the Flexi-UNet against the full UNet and under different architectural parametrizations:

Configuration	FID	Inference Time (s/image)
Full UNet, 2 levels	2.33	16.19
Flexi-UNet, 2 levels	2.26	15.38
Flexi-UNet, 4 levels	4.87	13.43

These results indicate that Flexi-UNet can match or slightly improve perceptual metrics while benefiting from reduced computational demand at low-to-intermediate levels.

Schedules that allocate more steps to higher resolutions achieve better FID but slow down training. Thus, Flexi-UNet exposes explicit control over fidelity-computation tradeoffs via the choice of resolution schedule.

6. Limitations and Prospective Extensions

Flexi-UNet as evaluated utilizes DDPM-style sampling and a standard ADM UNet backbone. Potential improvements include integration with more advanced samplers (e.g., DDIM, DPM-solvers), resolution curriculum strategies, and per-level tuning of hyperparameters (SNR schedules, time embeddings). Flexi-UNet’s design is compatible with alternative architectures such as Vision-Transformer-based backbones (DiT, U-ViT) for SSD in latent-space domains.

The reported framework explores only 2× downsampling, but the generalized SSD formulation allows straightforward replacement of the linear degradation with Gaussian blurs, anisotropic filters, or other operators. A natural direction is to exploit this flexibility for further FLOP/memory reduction or perceptual quality improvements.

A plausible implication is that Flexi-UNet's coarse-to-fine computation policy may inform future generative model design, particularly in domains where hierarchical information flow and efficient computation at multiple scales are crucial (Mukhopadhyay et al., 9 Mar 2026).

7. Context and Significance

Flexi-UNet provides a minimal yet effective modification to canonical diffusion model architectures, targeting the computational inefficiency in conventional approaches when processing highly degraded (noisy or low-resolution) states. By adapting computation to the meaningful spatial content at each timestep, it realizes substantial savings in compute and memory while preserving the capacity to yield high-fidelity samples. This framework demonstrates that generative processes need not uniformly allocate resources to all timesteps, especially when coarser representations suffice. The alignment of Flexi-UNet and Scale Space Diffusion stands as a principled instance of scale-adaptive generative modeling, with broad implications for scalable diffusion-based architectures (Mukhopadhyay et al., 9 Mar 2026).