---
title: Multi-Head cVAE for Disentangled Generation
url: https://www.emergentmind.com/topics/multi-head-conditional-variational-autoencoder-cvae
type: topic
---

# Multi-Head cVAE for Disentangled Generation

A Multi-Head Conditional Variational Autoencoder (cVAE) is a generative model architecture designed to decompose the latent representation into orthogonal components: a label-relevant code capturing structured, controllable information, and a label-irrelevant code capturing complementary, unsupervised variation. In the context of "Disentangling the Spatial Structure and Style in Conditional VAE" [1910.13062], this is achieved via a dual-headed latent space where one head encodes spatial structure or style associated with class labels, and the other head encodes class-independent factors. Each head is injected into the decoder with dedicated adaptive normalization mechanisms, enabling effective disentanglement of spatial structure and style in image generation.

## 1. Model Architecture

The multi-head cVAE consists of three major modules: a label-condition mapping network $f(\cdot)$ generating a label-relevant code $z_s$, an encoder producing a label-irrelevant latent code $z_u$, and a decoder conditioned on both $z_s$ and $z_u$ at every upsampling layer.

- **Label-Condition Mapping ($f(\cdot) \rightarrow z_s$):**
  - Input: a one-hot label $c \in \{0,1\}^{N_c}$.
  - Architecture: a multi-layer perceptron (3–4 fully-connected layers, width ≈ 512).
  - Output: embedding $z_s = f(c)$, which may be shaped as a spatial map $(H/k) \times (W/k) \times 1$ (if $c$ carries spatial information, e.g., pose) or a vector $1 \times 1 \times C$ (for categorical labels). Practical choices include $C=256$ and $k=4$ for $64 \times 64$ images (so $z_s$ shape is $16 \times 16 \times 1$).

- **Encoder ($\mathrm{Enc} \rightarrow z_u$):**
  - Input: image $x \in \mathbb{R}^{H \times W \times 3}$ (concatenated with label maps if necessary).
  - Architecture: strided convolutional blocks downsampling to either a $1 \times 1$ vector (style-posterior) or spatial map $(H/k) \times (W/k) \times 1$ (structure-posterior).
  - Latent outputs: mean $\mu(x, c)$ and std $\sigma(x, c)$ parameterizing $q_\phi(z_u \mid x, c) = \mathcal{N}(z_u; \mu(x, c), \text{diag}(\sigma^2(x, c)))$.

- **Decoder:**
  - Begins with a learned constant input.
  - Each upsampling block receives both $z_s$ (via SPADE) and $z_u$ (via AdaIN) to modulate activations.

## 2. Probabilistic Framework

Let $z_s$ denote the deterministic, label-relevant (label "head") code, and $z_u$ denote the stochastic, label-irrelevant (uncorrelated "head") code.

- **Priors:**
  - $p(z_u) = \mathcal{N}(z_u; 0, I)$ (isotropic Gaussian).
  - $p(z_s \mid c) = \delta(z_s - f(c))$ (deterministic).

- **Posteriors:**
  - $q_\phi(z_u \mid x, c) = \mathcal{N}(z_u; \mu(x,c), \text{diag}(\sigma^2(x,c)))$
  - $q(z_s \mid c) = \delta(z_s - f(c))$

- **Sampling:**
  - $z_u = \mu(x,c) + \sigma(x,c) \odot \epsilon$, $\epsilon \sim \mathcal{N}(0, I)$ (reparameterization)
  - $z_s = f(c)$ (deterministic)

- **ELBO Objective:**
  $$
  \mathcal{L}_{\text{ELBO}} = 
    \mathbb{E}_{q(z_u|x,c),\,q(z_s|c)}[\log p(x|z_u, z_s)] 
    - \mathrm{KL}[q(z_u|x,c)\|p(z_u)] 
    - \mathrm{KL}[q(z_s|c)\|p(z_s)]
  $$
  Given $q(z_s|c) = p(z_s|c)$, the last term vanishes. The likelihood is implemented as an $L_1$ or $L_2$ image reconstruction loss.

- **Adversarial Learning:**
  - Uses a cGAN-style hinge loss to sharpen outputs. A discriminator $D$ distinguishes between real and generated data, including cases with permuted labels or random $z_u$.

## 3. Adaptive Normalization in Decoding

At each decoder layer $l$, spatial and style codes modulate the activations via two normalization modules:

- **SPADE (label-relevant $z_s$):**
  - Produces $\gamma_s^{(l)}$, $\beta_s^{(l)}$ with spatial dimensions, via a small convolutional network upsampling $z_s$ to $H^{(l)} \times W^{(l)}$.

- **AdaIN (label-irrelevant $z_u$):**
  - Produces channel-wise $\gamma_u^{(l)}$, $\beta_u^{(l)}$ using an MLP applied to $z_u$ and broadcast spatially.

Given pre-activation $h^{(l)}$, normalization proceeds:
\[
\begin{align*}
\hat{h}_s^{(l)} &= \gamma_s^{(l)} \odot \frac{h^{(l)} - \mu_h^{(l)}}{\sigma_h^{(l)}} + \beta_s^{(l)} \\
\hat{h}_u^{(l)} &= \gamma_u^{(l)} \odot \frac{h^{(l)} - \mu_h^{(l)}}{\sigma_h^{(l)}} + \beta_u^{(l)}
\end{align*}
\]
The output features $\hat{h}_s^{(l)}$ and $\hat{h}_u^{(l)}$ are concatenated along the channel dimension and projected via a $1 \times 1$ convolution to restore channel size.

## 4. Implementation Configurations and Ablations

Key practical choices and architectural variants:

| Variant | $z_s$ injection | $z_u$ injection    |
|---------|-----------------|-------------------|
| S1      | AdaIN           | concat-input      |
| S2      | SPADE           | concat-input      |
| S3      | AdaIN           | AdaIN             |
| S4      | AdaIN           | SPADE             |
| Proposed| SPADE           | AdaIN             |

- Both $z_s$ and $z_u$ are dimensioned to 256, yielding in the structure code case $z_s$ of $16 \times 16 \times 1$ (for $64 \times 64$ images) and $z_u$ as $1 \times 1 \times 256$, or vice versa for style code scenarios.
- Encoder and decoder convolutional blocks follow channel progression $64 \rightarrow 128 \rightarrow 256 \rightarrow 512$.
- Default datasets: 3D-Chair ($64 \times 64$), FaceScrub ($128 \times 128$).
- Optimizer: Adam; learning rate and batch size are not fixed in the paper but typical settings are used (e.g., $lr=2 \times 10^{-4}$, $\beta_1=0.5$, $\beta_2=0.999$).

## 5. Quantitative and Qualitative Performance

Performance of the proposed disentangling design is demonstrated via experiments on 3D-Chair (label captures azimuth/viewpoint) and FaceScrub datasets (label as identity):

- **3D-Chair:**
  - Mutual Information $\mathcal{I}(z_u; c) = 3.750$ (lower is better; indicates improved disentanglement).
  - ResNet-50 classification accuracy at target azimuth: $0.623$.

- **FaceScrub:**
  - Identity-classification accuracy: $0.632$.
  - Fréchet Inception Distance (FID): $50.14$.

Qualitative results show that reconstructed or generated samples can enforce a target identity or viewpoint while preserving complementary factors such as style, pose, or expression.

## 6. Significance and Context

This design cleanly separates label-associated (structured) factors from unsupervised (residual) variation. By employing SPADE and AdaIN at every decoder stage—feeding the label-relevant and label-irrelevant codes respectively—conditional cVAE generation becomes modular, controllable, and suited for tasks demanding disentanglement. The approach enables, for example, faithful identity swapping in faces or view manipulation in 3D objects, where label information may or may not carry spatial meaning. This separation of signal pathways, along with the adversarial sharpness constraint, is shown to outperform simpler approaches in terms of both disentanglement metrics and visual fidelity [1910.13062].

Source: https://www.emergentmind.com/topics/multi-head-conditional-variational-autoencoder-cvae