---
title: 'SLIM-Diff: Compact Joint Diffusion for Imaging'
url: https://www.emergentmind.com/topics/slim-diff
type: topic
---

# SLIM-Diff: Compact Joint Diffusion for Imaging

SLIM-Diff denotes a compact joint image–mask diffusion framework designed for data-scarce medical imaging regimes, particularly for epilepsy-focused FLAIR MRI. Its core contributions are (i) a single shared-bottleneck U-Net enforcing coupling between anatomy and lesion geometry through a 2-channel representation, and (ii) explicit loss-geometry tuning via a tunable $L_p$ objective. Unlike standard large-scale diffusion architectures, SLIM-Diff is tailored for low-data settings and the simultaneous synthesis of both FLAIR images and corresponding lesion masks, targeting robust data augmentation and generative segmentation for rare disorders such as focal cortical dysplasia (FCD) [2602.03372]. 

## 1. Background and Motivation

Medical imaging datasets for rare neurological disorders, such as FCD II, are highly limited, typically yielding no more than $\sim$80 lesion-positive samples in public collections. This scarcity, combined with the subtle and highly-localized nature of lesion morphology on FLAIR MRI, triggers instability and overfitting in high-capacity generative models (e.g., canonical Stable Diffusion or LDM U-Nets with $\sim$860M parameters). The simultaneous generation of anatomically plausible images and spatially consistent lesion masks can enable task-specific data augmentation for segmentation and classification pipelines—but only if the generative process faithfully preserves both modalities' structure and does not merely memorize the limited training set.

SLIM-Diff addresses these requirements by adopting a tightly-coupled architecture (shared-bottleneck U-Net), reducing parameter count by more than an order of magnitude, and optimizing with a generalized loss geometry rather than canonical $L_2$ regression [2602.03372].

## 2. Mathematical Formulation

SLIM-Diff implements the discrete-time DDPM framework ($T=1000$ steps, cosine noise schedule) for a joint image–mask tensor $\mathbf{x}_0\in\mathbb{R}^{2\times H\times W}=[I,M]^\top$. The forward process adds Gaussian noise:
\[
q(\mathbf{x}_t | \mathbf{x}_0) = \mathcal{N}\left(\mathbf{x}_t;\,\sqrt{\bar\alpha_t}\,\mathbf{x}_0,\,(1-\bar\alpha_t)\mathbf{I}\right),\qquad\bar\alpha_t=\prod_{s=1}^t \alpha_s
\]
The reverse process, parameterized by U-Net $f_\theta$, targets one of three prediction parameterizations:
- $\epsilon$-prediction (standard, predicts noise)
- velocity $v_\theta = \sqrt{\bar\alpha_t}\,\epsilon_\theta - \sqrt{1-\bar\alpha_t}\,\mathbf{x}_0$
- direct $x_{0,\theta}$-prediction (clean image + mask)

The loss is a general $L_p$ norm over the prediction target:
\[
\mathcal{L} = \mathbb{E}_{t,\mathbf{x}_0,\epsilon} \|\mathrm{target} - f_\theta(\mathbf{x}_t, t, c)\|_p^p
\]
where $p$ is tunable ($p\in\{1.5,2.0,2.5\}$), with the specific choice impacting fidelity vs. mask morphology.

## 3. Shared-Bottleneck U-Net Architecture

The SLIM-Diff U-Net is strictly single-stream: both FLAIR image and binary lesion mask are concatenated as input. The encoder consists of four downsampling levels with channel widths $[64, 128, 256, 256]$, each with two residual blocks using GroupNorm (32 groups), with strided convolutions for spatial reduction. Multihead self-attention (32-dim heads) is included at the deepest two scales. A $20\times 20\times 256$ shared bottleneck compresses both modalities, imposing an information choke point to enforce representational coupling.

The decoder mirrors the encoder. Skip connections preserve spatial resolution, and a final $1\times1$ convolution predicts both modalities. Conditioning is supplied via:
- Axial + pathology indicator (60 discrete tokens: 30 axial bins $\times$ {control, lesion}), embedded via a learnable embedding and sinusoidal positional encoding.
- Timestep encoding (sinusoidal + 2-layer MLP).
- Both are injected into ResBlocks via FiLM-style bias modulation after the first convolution.

The total parameter count is $\sim$ 26.9M, a fraction of large-scale DMs.

## 4. Loss Geometry and Empirical Analysis

SLIM-Diff replaces the canonical $\epsilon$-prediction $L_2$ loss with explicit tuning over target and $p$:
- For $x_0$-prediction, $p=1.5$ (sub-quadratic) down-weights large residuals (e.g., hyperintense lesion pixels), improving global image realism (measured by KID, LPIPS).
- $p=2.0$ gives optimal mask boundary sharpness (lowest MMD-MF), as the uniform quadratic penalty preserves subtle geometry.
- $p=2.5$ over-penalizes outliers, degrading performance.
- Across all $p$, $x_0$-prediction significantly outperforms $\epsilon$-prediction and velocity-based targets in both image and mask metrics.

Table: Summary of best/typical configurations

| Target / $p$          | KID   | LPIPS | MMD-MF |
|----------------------|-------|-------|--------|
| $\epsilon$, $p=2.0$  | 0.432 | 0.821 | 15.06  |
| $x_0$, $p=1.5$       | 0.012 | 0.305 | 1.43   |
| $x_0$, $p=2.0$       | 0.034 | 0.310 | 0.95   |

Qualitative samples with $x_0/p=1.5$ exhibit realistic FLAIR contrast with coherent lesion masks; $x_0/p=2.0$ sharpens mask boundaries at the cost of slightly over-smoothed intensities.

No evidence of pixel-wise memorization was found; Kernel Maximum Mean Discrepancy tests confirmed distributional (not copy-based) generation.

## 5. Training and Sampling Procedures

The network is trained via AdamW (learning rate $10^{-4}$), cosine annealing, weight EMA ($0.999$ decay), and early stopping (validation patience 25 epochs). Preprocessing includes MNI registration, skull-stripping (ANTS SyN), N4 bias correction, and percentile normalization. Data is organized at the subject and slice level, with lesion oversampling to counteract class imbalance.

Sampling employs DDIM (300 steps, $\eta=0.2$), conditioned on desired axial depth and pathology indicator. Both training and inference require explicit condition tokens.

Pseudocode (Algorithm 1—Training):

```python
for epoch in 1...max_epochs:
    for (x0, c) in minibatch:
        t ~ Uniform(1, T)
        eps ~ N(0, I)
        xt = sqrt(bar_alpha_t)*x0 + sqrt(1-bar_alpha_t)*eps
        output = f_theta(xt, t, c)
        L = ||target(x0, eps, t) - output||_p^p
        # Backpropagation and parameter update
        # EMA update
    # Early stopping check
```

Pseudocode (Algorithm 2—Sampling):

```python
x_S ~ N(0, I)
for s in S...1:
    t = schedule(s)
    out = f_theta_bar(x_s, t, c)
    x_{s-1} = DDIM_update(x_s, out, t, eta)
return x0_image, x0_mask
```

## 6. Limitations and Applicability

SLIM-Diff operates on 2D slices, potentially introducing slight inter-slice inconsistencies. No direct comparison is provided against multi-stream and multi-stage joint generative frameworks (such as MedSegFactory or brainSPADE) under matched data conditions. Extension to full 3D or enforcing explicit slice-to-volume consistency is suggested as future work. For clinical deployment, care must be taken to evaluate robustness across lesion types, scanners, and domains; fine-tuning may be necessary.

Memory and computational requirements are modest, supporting real-time slice-level synthesis on standard GPUs.

## 7. Code and Reproducibility

The official implementation and pretrained models are available at https://github.com/MarioPasc/slim-diff. To reproduce, preprocessing (registration, skull stripping, bias correction), slice extraction, and balanced splitting are required. Training proceeds with configurable --target and --p_norm options; at convergence, the sample.py script enables generation under specified anatomical or pathological conditions.

In summary, SLIM-Diff establishes a methodological foundation for joint image–mask diffusion in data-scarce medical regimes, demonstrating that shared low-capacity architectures with explicit loss-geometry control enable robust synthesis of anatomically faithful images and masks while resisting memorization, and supporting downstream data augmentation for rare-disease imaging research [2602.03372].

Source: https://www.emergentmind.com/topics/slim-diff