technical_report_semantic.md

SemDisDiffAE — Technical Report

SemDisDiffAE (Semantically Disentangled Diffusion AutoEncoder) — a fast diffusion autoencoder with a 128-channel spatial bottleneck built on FCDM (Fully Convolutional Diffusion Model) blocks. The encoder uses a VP-parameterized diagonal Gaussian posterior (learned log-SNR output head), and the decoder reconstructs via single-step VP diffusion.

This checkpoint is trained with DINOv2 semantic alignment and variance expansion regularization. The name is a nod to DRA (Page et al., 2026) whose disentangled representation alignment approach we largely follow here.

Contents

1. Architecture
   • FCDM Block · Encoder · Decoder · AdaLN · PDG
2. Decoder VP Diffusion Parameterization
   • Forward Process · Log SNR · Schedule · X-Prediction · Sampling
3. Stochastic Posterior
   • VP Log-SNR Parameterization · Variance Expansion Loss · Posterior Mode
4. Semantic Alignment
5. Design Choices
   • Convolutional Architecture · Single-Stride Encoder · Diffusion Decoding · Skip Connection and PDG
6. Training
   • Loss Functions · Optimizer · Data
7. Model Configuration
8. Inference
9. Results

References:

• FCDM — Kwon et al., Reviving ConvNeXt for Efficient Convolutional Diffusion Models, arXiv:2603.09408, 2026.
• SiD2 — Hoogeboom et al., Simpler Diffusion (SiD2): 1.5 FID on ImageNet512 with pixel-space diffusion, arXiv:2410.19324, ICLR 2025.
• DiTo — Yin et al., Diffusion Autoencoders are Scalable Image Tokenizers, arXiv:2501.18593, 2025.
• DiCo — Ai et al., DiCo: Revitalizing ConvNets for Scalable and Efficient Diffusion Modeling, arXiv:2505.11196, 2025.
• ConvNeXt V2 — Woo et al., ConvNeXt V2: Co-designing and Scaling ConvNets with Masked Autoencoders, arXiv:2301.00808, CVPR 2023.
• Z-image — Cai et al., Z-Image: An Efficient Image Generation Foundation Model with Single-Stream Diffusion Transformer, arXiv:2511.22699, 2025.
• SPRINT — Park et al., Sprint: Sparse-Dense Residual Fusion for Efficient Diffusion Transformers, arXiv:2510.21986, 2025.
• DINOv2 — Oquab et al., DINOv2: Learning Robust Visual Features without Supervision, arXiv:2304.07193, 2023. Register variant: Darcet et al., Vision Transformers Need Registers, arXiv:2309.16588, ICLR 2024.
• iREPA — Singh et al., What matters for Representation Alignment: Global Information or Spatial Structure?, arXiv:2512.10794, 2025.
• DRA — Page et al., Boosting Latent Diffusion Models via Disentangled Representation Alignment, arXiv:2601.05823, 2026.
• VEL — Li et al., Taming Sampling Perturbations with Variance Expansion Loss for Latent Diffusion Models, arXiv:2603.21085, 2026.
• iRDiffAE — data-archetype/irdiffae-v1 — predecessor model using DiCo blocks.

---

1. Architecture

1.1 FCDM Block

SemDisDiffAE uses FCDM blocks — ConvNeXt-style convolutional blocks adapted for diffusion models (Li et al., 2026). Each block follows a single unified residual path:

x ──► DWConv 7×7 ──► RMSNorm ──► [Scale] ──► Conv 1×1 ──► GELU ──► GRN ──► Conv 1×1 ──► [Gate] ──► + ──► out
│                                                                                                    ▲
└────────────────────────────────────────────────────────────────────────────────────────────────────────┘

This differs from DiCo blocks (used in the predecessor iRDiffAE) which use two separate residual paths (conv + MLP) with Compact Channel Attention (CCA). FCDM consolidates into one path, replacing CCA with Global Response Normalization (GRN).

Key components:

• Depthwise convolution (7×7, groups=channels): spatial mixing without cross-channel interaction. The depthwise conv output feeds directly into RMSNorm (no intermediate activation).

• RMSNorm (non-affine, per-channel): normalizes activations before the pointwise MLP, replacing LayerNorm used in standard ConvNeXt.

• Global Response Normalization (GRN) (from ConvNeXt V2, Woo et al. 2023): applied between the two pointwise convolutions. GRN computes per-channel L2 norms across the spatial dimensions and normalizes by the cross-channel mean:

  g = ||x||_2 over (H, W)
  n = g / mean(g over channels)
  GRN(x) = gamma * (x * n) + beta + x
  

  This encourages feature diversity across channels and prevents channel collapse during training.

• Scale+Gate modulation (decoder only): FCDM blocks use a 2-way modulation (scale, gate) from the timestep embedding, in contrast to DiCo's 4-way (shift_conv, gate_conv, shift_mlp, gate_mlp). Scale is applied after RMSNorm: h = h * (1 + scale). Gate is applied to the residual: out = x + gate * h. The gate is used raw (no tanh activation), giving unbounded gating — this differs from DiCo which applies tanh to constrain the gate to [-1, 1].

• Layer Scale (encoder only): for unconditioned encoder blocks, a learnable per-channel scale (initialized to 1e-3) gates the residual for near-identity initialization, following ConvNeXt.

1.2 Encoder

The encoder uses a single spatial stride (via PixelUnshuffle at the input) followed by FCDM blocks at constant spatial resolution, then a bottleneck projection that outputs both the posterior mean and per-element log-SNR:

Image [B, 3, H, W]
  ──► PixelUnshuffle(p=16) + Conv 1×1 (3·16² → 896)     [Patchify]
  ──► 4 × FCDMBlock (unconditioned, layer-scale gated)
  ──► Conv 1×1 (896 → 256)                               [Bottleneck projection]
  ──► Split → mean [B, 128, h, w] + logsnr [B, 128, h, w]
  ──► α(logsnr) · mean                                    [Posterior mode]

The single-stride design ensures all encoder blocks see the full spatial resolution and full channel width simultaneously. The information bottleneck is imposed only at the very end, where a single linear projection selects which channels to retain. See Section 4.2 for the rationale.

Note: This checkpoint uses bottleneck_norm_mode=disabled, so no post-bottleneck RMSNorm is applied to the mean branch. The posterior mode output is simply α · μ where α = √σ(λ).

1.3 Decoder

The decoder predicts x̂₀ from noisy input x_t, conditioned on encoder latents z and timestep t:

Noised image x_t [B, 3, H, W]
  ──► PixelUnshuffle(p=16) + Conv 1×1 (3·16² → 896)     [Patchify]
  ──► Concatenate with Conv 1×1(latents, 128 → 896)      [Latent fusion]
  ──► Conv 1×1 (2·896 → 896)
  ──► 2 × FCDMBlock (AdaLN conditioned)                   [Start blocks]
  ──► 4 × FCDMBlock (AdaLN conditioned)                   [Middle blocks]
  ──► Concat(start_out, middle_out) + Conv 1×1            [Skip fusion]
  ──► 2 × FCDMBlock (AdaLN conditioned)                   [End blocks]
  ──► Conv 1×1 (896 → 3·16²) + PixelShuffle(16)          [Unpatchify]
  ──► x̂₀ prediction [B, 3, H, W]

The skip-concat topology with 2+4+2 blocks is inspired by SPRINT's sparse-dense residual fusion (Park et al., 2025). See Section 4.4 for the design rationale.

1.4 AdaLN: Shared Base + Low-Rank Deltas

Timestep conditioning follows the Z-image style AdaLN (Cai et al., 2025): a shared base projection plus a low-rank delta per layer.

A single base projector is shared across all 8 decoder layers, and each layer adds a low-rank correction:

m_i = Base(SiLU(cond)) + Δ_i(SiLU(cond))

where Base: ℝ^D → ℝ^{2D} is a linear projection (zero-initialized) and Δ_i: ℝ^D → ℝ^r → ℝ^{2D} is a low-rank factorization with rank r = 128 (zero-initialized up-projection).

The packed modulation m_i ∈ ℝ^{B × 2D} is split into (scale, gate) which modulate the FCDM block (no shift term):

ĥ = RMSNorm(x) · (1 + scale)
x ← x + gate · f(ĥ)

1.5 Path-Drop Guidance (PDG)

At inference, optional PDG sharpens reconstructions by exploiting the skip-concat structure — a classifier-free guidance analogue that does not require training with conditioning dropout:

1. Conditional pass: run all blocks normally → x̂₀^cond
2. Unconditional pass: replace the middle block output with a learned mask feature m ∈ ℝ^{1×D×1×1} (initialized to zero), effectively dropping the deep processing path → x̂₀^uncond
3. Guided prediction: x̂₀ = x̂₀^uncond + s · (x̂₀^cond - x̂₀^uncond)

where s is the guidance strength.

For PSNR-optimal reconstruction, PDG is disabled (1 NFE). For perceptual sharpening, use 10 steps with PDG strength 2.0. Note that PDG is primarily useful for more compressed bottlenecks (e.g. 32 or 64 channels) and is rarely necessary for 128-channel models where reconstruction quality is already high.

---

2. Decoder VP Diffusion Parameterization

The decoder uses the variance-preserving (VP) diffusion framework from SiD2 with an x-prediction objective.

2.1 Forward Process

Given a clean image $x_{0}x\_0$, the forward process constructs a noisy sample at continuous time $t\in [0,1]t\; \backslash in\; [0,\; 1]$:

$$x_t={\alpha }_t{x}_0+{\sigma }_t\epsilon ,\epsilon \sim \mathcal{N}(0,s^{2}I)x\_t\; =\; \backslash alpha\_t\; \backslash ,\; x\_0\; +\; \backslash sigma\_t\; \backslash ,\; \backslash varepsilon,\; \backslash qquad\; \backslash varepsilon\; \backslash sim\; \backslash mathcal\{N\}(0,\; s^2\; I)$$

where $s=0.558s\; =\; 0.558$ is the pixel-space noise standard deviation (estimated from the dataset image distribution) and the VP constraint holds: ${\alpha }_t^2+{\sigma }_t^2=1\backslash alpha\_t^2\; +\; \backslash sigma\_t^2\; =\; 1$.

2.2 Log Signal-to-Noise Ratio

The schedule is parameterized through the log signal-to-noise ratio:

$${\lambda }_t=\log \frac{{\alpha }_t^2}{{\sigma }_t^2}\backslash lambda\_t\; =\; \backslash log\; \backslash frac\{\backslash alpha\_t^2\}\{\backslash sigma\_t^2\}$$

which monotonically decreases as $t\to 1t\; \backslash to\; 1$ (pure noise). From ${\lambda }_t\backslash lambda\_t$ we recover ${\alpha }_t\backslash alpha\_t$ and ${\sigma }_t\backslash sigma\_t$ via the sigmoid function:

$${\alpha }_t=\sqrt{\sigma ({\lambda }_t)},{\sigma }_t=\sqrt{\sigma (-{\lambda }_t)}\backslash alpha\_t\; =\; \backslash sqrt\{\backslash sigma(\backslash lambda\_t)\},\; \backslash qquad\; \backslash sigma\_t\; =\; \backslash sqrt\{\backslash sigma(-\backslash lambda\_t)\}$$

2.3 Cosine-Interpolated Schedule

Following SiD2, the logSNR schedule uses cosine interpolation:

$$\lambda (t)=-2\log \tan (a\cdot t+b)\backslash lambda(t)\; =\; -2\; \backslash log\; \backslash tan(a\; \backslash cdot\; t\; +\; b)$$

where $aa$ and $bb$ are computed to satisfy the boundary conditions $\lambda (0)={\lambda }_{\text{max}}=10\backslash lambda(0)\; =\; \backslash lambda\_\backslash text\{max\}\; =\; 10$ and $\lambda (1)={\lambda }_{\text{min}}=-10\backslash lambda(1)\; =\; \backslash lambda\_\backslash text\{min\}\; =\; -10$.

2.4 X-Prediction Objective

The model predicts the clean image ${\widehat{x}}_0=f_{\theta }(x_t,t,z)\backslash hat\{x\}\_0\; =\; f\_\backslash theta(x\_t,\; t,\; z)$ conditioned on encoder latents $zz$.

Schedule-invariant loss. Following SiD2, the training loss is defined as an integral over logSNR $\lambda \backslash lambda$, making it invariant to the choice of noise schedule. Since timesteps are sampled uniformly $t\sim \mathcal{U}(0,1)t\; \backslash sim\; \backslash mathcal\{U\}(0,1)$, the change of variable introduces a Jacobian factor:

$$\mathcal{L}={\mathbb{E}}_{t\sim \mathcal{U}(0,1)}[(-\frac{d\lambda }{dt})\cdot w(\lambda (t))\cdot \mathrm{\parallel }{x}_0-{\widehat{x}}_0{\mathrm{\parallel }}^2]\backslash mathcal\{L\}\; =\; \backslash mathbb\{E\}\_\{t\; \backslash sim\; \backslash mathcal\{U\}(0,1)\}\; \backslash left[\; \backslash left(-\backslash frac\{d\backslash lambda\}\{dt\}\backslash right)\; \backslash cdot\; w(\backslash lambda(t))\; \backslash cdot\; \backslash |\; x\_0\; -\; \backslash hat\{x\}\_0\; \backslash |^2\; \backslash right]$$

Sigmoid weighting. The weighting function uses a sigmoid centered at bias $b=-2.0b\; =\; -2.0$, converting from $\epsilon \backslash varepsilon$-prediction to $xx$-prediction form:

$$\text{weight}(t)=-\frac{1}{2}\frac{d\lambda }{dt}\cdot e^b\cdot \sigma (\lambda (t)-b)\backslash text\{weight\}(t)\; =\; -\backslash frac\{1\}\{2\}\; \backslash frac\{d\backslash lambda\}\{dt\}\; \backslash cdot\; e^b\; \backslash cdot\; \backslash sigma(\backslash lambda(t)\; -\; b)$$

2.5 Sampling

Decoding uses DDIM by default. With 1 NFE (default), the model runs a single evaluation at t_start ≈ 1 (near pure noise) and directly outputs the x₀ prediction. This is equivalent to a denoising autoencoder that maps σ_start · noise → x̂₀ conditioned on encoder latents.

DPM++2M is also supported as an alternative sampler, using a half-lambda exponential integrator for faster convergence with more steps.

---

3. Stochastic Posterior

3.1 VP Log-SNR Parameterization

Instead of a KL-divergence penalty on a Gaussian encoder, SemDisDiffAE parameterizes the bottleneck posterior using the VP interpolation convention. This approach uses a VP-style noise interpolation in the encoder bottleneck as an alternative to the traditional VAE KL penalty.

The encoder outputs two sets of 128 channels:

• $\mu \backslash mu$ — the clean signal (posterior mean)
• $\lambda \backslash lambda$ — per-element log signal-to-noise ratio

The posterior distribution is:

$$z=\alpha (\lambda )\mu +\sigma (\lambda )\epsilon ,\epsilon \sim \mathcal{N}(0,I)z\; =\; \backslash alpha(\backslash lambda)\; \backslash ,\; \backslash mu\; +\; \backslash sigma(\backslash lambda)\; \backslash ,\; \backslash varepsilon,\; \backslash qquad\; \backslash varepsilon\; \backslash sim\; \backslash mathcal\{N\}(0,\; I)$$

where $\alpha =\sqrt{\sigma (\lambda )}\backslash alpha\; =\; \backslash sqrt\{\backslash sigma(\backslash lambda)\}$ and $\sigma =\sqrt{\sigma (-\lambda )}\backslash sigma\; =\; \backslash sqrt\{\backslash sigma(-\backslash lambda)\}$ (sigmoid parameterization). This is equivalent to a Gaussian with mean $\alpha \mu \backslash alpha\; \backslash mu$ and variance ${\sigma }^2\backslash sigma^2$.

Using a VP interpolation rather than simple additive noise decouples token scale from stochasticity. With additive noise (\(z = \mu + \sigma\varepsilon\)), the encoder faces gradient pressure to scale latents up to counter the noise — the SNR depends on the magnitude of $\mu \backslash mu$. The VP formulation (\(z = \alpha\mu + \sigma\varepsilon\) with ${\alpha }^2+{\sigma }^2=1\backslash alpha^2\; +\; \backslash sigma^2\; =\; 1$) removes this coupling: the noise level is controlled entirely by the predicted log-SNR, independent of the latent magnitude.

3.2 Variance Expansion Loss

To prevent posterior collapse (where the encoder learns to set σ → 0 and ignore the stochastic component entirely), we adopt a variance expansion loss inspired by VEL (Li et al., 2026, arXiv:2603.21085):

$${\mathcal{L}}_{\text{var}}=-\mathrm{mean}\text{\hspace{0.17em}⁣}(\log ({\sigma }^2+\delta ))\backslash mathcal\{L\}\_\backslash text\{var\}\; =\; -\backslash operatorname\{mean\}\backslash !\backslash bigl(\backslash log(\backslash sigma^2\; +\; \backslash delta)\backslash bigr)$$

where ${\sigma }^2\backslash sigma^2$ is the posterior variance derived from the predicted log-SNR and $\delta ={10}^{-6}\backslash delta\; =\; 10^\{-6\}$ for numerical stability. This loss encourages non-zero posterior variance by penalizing small ${\sigma }^2\backslash sigma^2$.

VEL proposes the form $1\mathrm{/}({\sigma }^2+\delta )1/(\backslash sigma^2\; +\; \backslash delta)$ for variance expansion. We found this to be too aggressive — the $1\mathrm{/}{\sigma }^{2}1/\backslash sigma^2$ gradient pushes variance up very rapidly, leading to excessive high-frequency noise in the latent space. We use the $-\log ({\sigma }^2+\delta )-\backslash log(\backslash sigma^2\; +\; \backslash delta)$ form instead, which provides a gentler, logarithmic penalty that stabilizes training.

For this checkpoint: the variance expansion loss is active with weight 1e-5.

Key finding: latent spectral structure matters for downstream diffusion.

Reconstruction quality is not very sensitive to the posterior noise level — good PSNR is achievable even with log-SNR as low as -2. However, the posterior noise level has a strong effect on the spatial frequency content of the latent space. When variance expansion is too aggressive, the latent space develops excessive high-frequency content; when it is too weak or absent, latents become overly smooth.

We found empirically that downstream diffusion models converge best when the latent space has a radial power spectral density (PSD) decay exponent of approximately 1.5 — deviating significantly in either direction (too smooth or too high-frequency) consistently yields worse downstream training convergence. We monitor this metric during training validation to guide the variance expansion weight.

The weight of 1e-5 for this checkpoint was chosen to target this spectral sweet spot.

3.3 Posterior Mode for Inference

At inference, the encoder returns the posterior mode: z = α(λ) · μ. For this checkpoint, the posterior log-SNR is very high (posterior variance is negligible), so sampling and mode are nearly identical.

The encode_posterior() method is available for users who need the full posterior distribution.

---

4. Semantic Alignment

This checkpoint uses semantic alignment to encourage semantically structured latent representations. The approach is inspired by DRA (Page et al., 2026, arXiv:2601.05823) which aligns autoencoder latents with frozen vision encoder features. Our implementation differs in the projection architecture and noise schedule.

4.1 Teacher

A frozen DINOv2-S with registers (timm: vit_small_patch16_dinov3.lvd_1689m, 384-dim patch tokens) provides the target spatial semantic features.

4.2 Projection Head

The student projection head maps noisy encoder latents to the teacher's token space. It consists of:

Noisy latents z_noisy ∈ ℝ^{B×128×h×w}
  ──► Conv 1×1 (128 → 384)                     [Channel projection]
  ──► Flatten to tokens [B, T, 384]
  ──► DiT transformer block                     [Single block, 6 heads × 64 dim]
      (self-attention with axial RoPE 2D + AdaLN conditioned on τ)
  ──► RMSNorm
  ──► student tokens ∈ ℝ^{B×T×384}

The DiT block uses standard multi-head self-attention with 2D axial rotary position embeddings (RoPE) and AdaLN-Zero timestep conditioning. This gives the projection head global spatial reasoning — important for matching the teacher's self-attention-based representations — while the main encoder/decoder remain purely convolutional.

4.3 Noisy Alignment

Unlike standard representation alignment which operates on clean latents, we align noisy latent versions. The noise level $\tau \backslash tau$ is sampled from a $\text{Beta}(2,2)\backslash text\{Beta\}(2,2)$ distribution (concentrated around $\tau =0.5\backslash tau\; =\; 0.5$) using flow matching linear interpolation:

$$z_{\text{noisy}}=(1-\tau )z+\tau \epsilon ,\epsilon \sim \mathcal{N}(0,I),\tau \sim \text{Beta}(2,2)z\_\backslash text\{noisy\}\; =\; (1\; -\; \backslash tau)\; \backslash ,\; z\; +\; \backslash tau\; \backslash ,\; \backslash varepsilon,\; \backslash qquad\; \backslash varepsilon\; \backslash sim\; \backslash mathcal\{N\}(0,\; I),\; \backslash quad\; \backslash tau\; \backslash sim\; \backslash text\{Beta\}(2,\; 2)$$

The projection head receives both the noisy latents and the noise level $\tau \backslash tau$ (via its AdaLN conditioning). This trains the head to extract semantic information even from partially corrupted latents, improving robustness for downstream diffusion models which operate on noised latent inputs.

4.4 Training Details

The alignment loss is the mean negative cosine similarity between student and teacher tokens, weighted at 0.01 throughout training. The student projection head operates on all 128 bottleneck channels, unlike the predecessor iRDiffAE which aligned only the first 64 of 128 channels.

Note that the projection head is a training-only component — it is not included in the exported model weights.

---

5. Design Choices

5.1 Convolutional Architecture

SemDisDiffAE uses a fully convolutional architecture rather than a vision transformer. For an autoencoder whose goal is faithful pixel-level reconstruction (not global semantic understanding), convolutions offer:

• Resolution generalization. Convolutions operate on local patches and generalize naturally to arbitrary image dimensions without interpolating position embeddings or suffering attention distribution shift.
• Translation invariance. Weight sharing across spatial positions is well matched to reconstruction, where the same local patterns (edges, textures) conditioned on the low-frequency latent recur throughout the image.
• Locality. Reconstruction quality depends on preserving fine spatial detail. Convolutions are inherently local operators, avoiding the quadratic cost of global attention while focusing computation where it matters most.

5.2 Single-Stride Encoder with Final Bottleneck

The encoder uses a single spatial stride (PixelUnshuffle at the input) followed by blocks at constant spatial resolution, then a final 1×1 convolution to project to the bottleneck. This differs from classical VAE encoders that use progressive downsampling with channel expansion at each stage.

The single-stride design ensures that all encoder blocks see the full spatial resolution and full channel width simultaneously. The information bottleneck is imposed only at the very end, where a single linear projection selects which channels to retain.

5.3 Diffusion Decoding

The main advantage of diffusion decoding over the standard GAN + LPIPS approach is simplicity and speed of experimentation. The training objective is a straightforward weighted MSE — no discriminator, no LPIPS perceptual loss, no delicate adversarial balancing. This makes it very fast to train and easy to iterate on — typically a few hours on a single GPU is sufficient. This checkpoint was trained for 251k steps. By contrast, GAN + LPIPS-based VAEs require many days of large-GPU time and are notoriously difficult to stabilize from scratch.

This simplicity enables rapid experimentation with latent space shaping to get it as diffusion-friendly as possible, while still achieving excellent reconstruction quality.

5.4 Skip Connection and Path-Drop Guidance

The decoder's start → middle → skip-fuse → end architecture is inspired by SPRINT's sparse-dense residual fusion (Park et al., 2025). The design serves three purposes:

1. Regularization. The skip path ensures that even if the middle blocks are dropped or poorly conditioned, the end blocks still receive meaningful features from the start blocks.
2. High-frequency preservation. The start blocks (which see the input most directly) pass fine detail through the skip to the end blocks.
3. Path-Drop Guidance. At inference, replacing the middle block output with a learned mask feature creates an "unconditional" prediction that preserves the skip path but drops the deep processing. Interpolating between conditional and unconditional predictions (as in classifier-free guidance) sharpens the output without requiring training-time dropout.

---

6. Training

6.1 Loss Functions

The total training loss is:

$${\mathcal{L}}_{\text{total}}={\mathcal{L}}_{\text{recon}}+0.01\cdot {\mathcal{L}}_{\text{semantic}}+{10}^{-4}\cdot {\mathcal{L}}_{\text{scale}}+{10}^{-5}\cdot {\mathcal{L}}_{\text{var}}\backslash mathcal\{L\}\_\backslash text\{total\}\; =\; \backslash mathcal\{L\}\_\backslash text\{recon\}\; +\; 0.01\; \backslash cdot\; \backslash mathcal\{L\}\_\backslash text\{semantic\}\; +\; 10^\{-4\}\; \backslash cdot\; \backslash mathcal\{L\}\_\backslash text\{scale\}\; +\; 10^\{-5\}\; \backslash cdot\; \backslash mathcal\{L\}\_\backslash text\{var\}$$

Loss	Weight	Description
${\mathcal{L}}_{\text{recon}}\backslash mathcal\{L\}\_\backslash text\{recon\}$	1.0	SiD2 sigmoid-weighted x-prediction MSE (\(b = -2.0\)). Per-pixel $(hat{x}_0-x_0{)}^2(\backslash \backslash hat\{x\}\_0\; -\; x\_0)^2$ averaged over (C, H, W) per sample, multiplied by $w(t)=-\frac{1}{2}\frac{d\lambda }{dt}{e}^b\sigma (\lambda -b)w(t)\; =\; -\backslash tfrac\{1\}\{2\}\; \backslash tfrac\{d\backslash lambda\}\{dt\}\; e^b\; \backslash sigma(\backslash lambda\; -\; b)$, then averaged over the batch
${\mathcal{L}}_{\text{semantic}}\backslash mathcal\{L\}\_\backslash text\{semantic\}$	0.01	Per-token $1-\cos (\text{student},\text{teacher})1\; -\; \backslash cos(\backslash text\{student\},\; \backslash text\{teacher\})$ averaged over all tokens and batch (see §4)
${\mathcal{L}}_{\text{scale}}\backslash mathcal\{L\}\_\backslash text\{scale\}$	0.0001	Per-channel variance ${\text{var}}_c\backslash text\{var\}\_c$ estimated over (B, H, W), then $(\log ({\text{var}}_c+\epsilon )-\log (\text{target}){)}^2(\backslash log(\backslash text\{var\}\_c\; +\; \backslash varepsilon)\; -\; \backslash log(\backslash text\{target\}))^2$ averaged over channels. Target variance = 1.0
${\mathcal{L}}_{\text{var}}\backslash mathcal\{L\}\_\backslash text\{var\}$	1e-5	Per-element $-\log ({\sigma }^2+\delta )-\backslash log(\backslash sigma^2\; +\; \backslash delta)$ where ${\sigma }^2\backslash sigma^2$ is the posterior variance, averaged over all dims (B, C, H, W). See §3.2

Note on loss scales: The decoder reconstruction loss has a small effective magnitude due to the SiD2 VP x-prediction weighting (the Jacobian $d\lambda \mathrm{/}dtd\backslash lambda/dt$ and sigmoid weighting compress the per-sample loss scale). As a result, all auxiliary loss weights must be kept correspondingly small to avoid dominating the reconstruction objective.

6.2 Optimizer and Hyperparameters

Parameter	Value
Optimizer	AdamW (β₁=0.9, β₂=0.99)
Learning rate	1e-4 (constant after warmup)
Weight decay	0.0
Warmup steps	2,000
Gradient clip	1.0 (max norm)
Precision	AMP bfloat16 (FP32 master weights, TF32 matmul)
EMA decay	0.9995 (updated every step)
Batch size	128
Timestep sampling	Uniform with SiD2 logSNR shift -1.0
Compilation	torch.compile enabled
Training steps	251k
Hardware	Single GPU

Convergence is fast — training is stopped when the training loss starts plateauing, which typically occurs within a few hours on a single GPU.

6.3 Data

Training uses ~5M images at various resolutions: mostly photographs, with a significant proportion of illustrations and text-heavy images (documents, screenshots, book covers, diagrams) to encourage crisp line and edge reconstruction. Images are loaded via two strategies in a 50/50 mix:

• Full-image downsampling: images are bucketed by aspect ratio and downsampled to ~256² resolution (preserving aspect ratio).
• Random 256×256 crops: deterministic patches extracted from images stored at ≥512px resolution.

This mixed strategy exposes the model to both global scene composition (via downsampled full images) and fine local detail (via crops from higher-resolution sources).

---

7. Model Configuration

Parameter	Value
Patch size	16
Model dimension	896
Encoder depth	4 blocks
Decoder depth	8 blocks (2 start + 4 middle + 2 end)
Bottleneck dimension	128 channels
Spatial compression	16× (H/16 × W/16)
Total compression	6.0× (3·256 / 128)
MLP ratio	4.0
Depthwise kernel	7×7
AdaLN per-block delta rank	128
Block type	FCDM (ConvNeXt + GRN + scale/gate AdaLN)
Posterior	Diagonal Gaussian (VP log-SNR), variance expansion weight 1e-5
Bottleneck norm	Disabled
λ_min, λ_max	-10, +10
Sigmoid bias b	-2.0
Pixel noise std s	0.558
Parameters	88.8M

---

8. Inference

Recommended Settings

Use case	Steps (NFE)	PDG	Sampler	Notes
PSNR-optimal	1	off	DDIM	Default. Fastest.
Perceptual	10	on (2.0)	DDIM	Sharper details, ~15× slower (PDG skips middle blocks)

Usage

from fcdm_diffae import FCDMDiffAE, FCDMDiffAEInferenceConfig

# Load model
model = FCDMDiffAE.from_pretrained("data-archetype/semdisdiffae", device="cuda")

# Encode (returns posterior mode by default)
latents = model.encode(images)  # [B,3,H,W] → [B,128,H/16,W/16]

# Decode (1 step)
recon = model.decode(latents, height=H, width=W)

# Full posterior access
posterior = model.encode_posterior(images)
print(posterior.mean.shape, posterior.logsnr.shape)
z_sampled = posterior.sample()

---

Citation

@misc{semdisdiffae,
  title   = {SemDisDiffAE: A Semantically Disentangled Diffusion Autoencoder},
  author  = {data-archetype},
  email   = {data-archetype@proton.me},
  year    = {2026},
  month   = apr,
  url     = {https://huggingface.co/data-archetype/semdisdiffae},
}

---

9. Results

Reconstruction quality evaluated on a curated set of test images covering photographs, book covers, and documents.

7.1 Interactive Viewer

Open full-resolution comparison viewer — side-by-side reconstructions, RGB deltas, and latent PCA with adjustable image size.

7.2 Inference Settings

Setting	Value
Sampler	ddim
Steps	1
Schedule	linear
Seed	42
PDG	no_path_dropg
Batch size (timing)	4

All models run in bfloat16. Timings measured on an NVIDIA RTX Pro 6000 (Blackwell).

7.3 Global Metrics

Metric	semdisdiffae (1 step)	Flux.2 VAE
Avg PSNR (dB)	35.78	34.16
Avg encode (ms/image)	2.5	46.1
Avg decode (ms/image)	5.5	91.8

7.4 Per-Image PSNR (dB)

Image	semdisdiffae (1 step)	Flux.2 VAE
p640x1536:94623	35.44	33.50
p640x1536:94624	31.33	30.03
p640x1536:94625	35.05	33.98
p640x1536:94626	33.21	31.53
p640x1536:94627	32.54	30.53
p640x1536:94628	29.80	28.88
p960x1024:216264	46.37	45.39
p960x1024:216265	29.70	27.80
p960x1024:216266	47.15	46.20
p960x1024:216267	40.99	39.23
p960x1024:216268	38.47	36.13
p960x1024:216269	32.74	30.24
p960x1024:216270	36.23	34.18
p960x1024:216271	44.41	42.18
p704x1472:94699	43.80	41.79
p704x1472:94700	32.83	32.08
p704x1472:94701	39.00	37.90
p704x1472:94702	34.52	32.50
p704x1472:94703	32.81	31.35
p704x1472:94704	33.38	31.84
p704x1472:94705	39.70	37.44
p704x1472:94706	35.12	33.66
r256_p1344x704:15577	31.02	29.98
r256_p1344x704:15578	32.38	30.79
r256_p1344x704:15579	33.27	31.83
r256_p1344x704:15580	37.84	36.03
r256_p1344x704:15581	38.57	36.94
r256_p1344x704:15582	33.41	32.10
r256_p1344x704:15583	36.67	34.54
r256_p1344x704:15584	33.23	31.76
r256_p896x1152:144131	35.30	33.60
r256_p896x1152:144132	36.99	35.32
r256_p896x1152:144133	39.69	37.33
r256_p896x1152:144134	36.01	34.47
r256_p896x1152:144135	31.20	29.87
r256_p896x1152:144136	37.51	35.68
r256_p896x1152:144137	33.83	32.86
r256_p896x1152:144138	27.39	25.63
VAE_accuracy_test_image	36.64	35.25