Symmetry-Conditioned VAE

• Symmetry-Conditioned VAE is an unsupervised model that leverages learnable group symmetries to align latent factors without requiring factor labels.
• It utilizes an extended ELBO combined with algebraic and equivariance regularizations to enforce disentangled, axis-aligned latent representations.
• The approach employs a learnable symmetry codebook and composite group operators to achieve multi-factor disentanglement, evaluated via the m-FVMₖ metric.

A Symmetry-Conditioned Variational Autoencoder (CVAE), as defined in the Composite Factor-Aligned Symmetry Learning (CFASL) framework, is a class of unsupervised generative latent variable models in which disentanglement emerges from explicit, learnable group symmetries acting on the latent space. Unlike prior approaches that require factor labels or known generative factors, the CFASL methodology enables a VAE to discover and align latent dimensions with independent generative symmetries by integrating a suite of algebraic regularization and equivariance conditions, operationalized entirely through a learnable symmetry codebook and composite group actions inferred from pairs of data samples (Jung et al., 2024).

1. Loss Augmentation: ELBO with Symmetry and Equivariance Constraints

The foundational loss function of the CVAE is an extended evidence lower bound (ELBO), augmented with several families of regularization and alignment losses. The total loss is: $$\mathcal{L}(\phi,\theta) = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x\mid z)] - \beta\,D_{\mathrm{KL}}(q_\phi(z|x)\,\|\,p(z)) + \lambda_c\mathcal{L}_c + \lambda_{pl}\mathcal{L}_{pl} + \lambda_{pd}\mathcal{L}_{pd} + \lambda_s\mathcal{L}_s + \lambda_p\mathcal{L}_p + \lambda_{ee}\mathcal{L}_{ee} + \lambda_{de}\mathcal{L}_{de}$$ where each $\lambda$ is a hyper-parameter and the terms $$\mathcal{L}_c, \mathcal{L}_{pl}, \mathcal{L}_{pd}, \mathcal{L}_s, \mathcal{L}_p$$ encode algebraic and statistical constraints aligned with symmetry discovery. $\mathcal{L}_{ee}$ and $\mathcal{L}_{de}$ enforce group equivariance in the encoder and decoder, respectively. This composite loss systematically injects inductive bias for the emergence of factor-aligned, axis-aligned, and group-structured latent representations.

2. Learnable Symmetry Codebook

The central algebraic structure is the symmetry codebook $S$. For latent dimension $d$, a library $$\mathfrak{B} = \{B_1,\dots,B_B\} \subset \mathbb{R}^{d\times d}$$ forms a basis for the relevant Lie algebra. The codebook $S$ is partitioned into $K$ sections, one per postulated generative factor: $$S = \{\mathcal{G}^1, \mathcal{G}^2, \dots, \mathcal{G}^K\}$$ with each $$\mathcal{G}^i = \{\mathfrak{g}^i_1, \dots, \mathfrak{g}^i_M\}$$, and $\mathfrak{g}^i_j$ parameterized as trainable combinations of basis elements. Group elements are recovered via the matrix exponential $g^i_j = \exp(\mathfrak{g}^i_j)$. Each section mediates transformations corresponding to a candidate latent factor, enabling unsupervised discovery and alignment of symmetries with latent dimensions through explicit parametrization and subsequent algebraic regularization.

3. Composite Symmetry Operators and Factor Selection

To capture multi-factor generative changes, CFASL infers composite group elements from paired data. Given images $(x^1,x^2)$, their posterior parameters $(\mu^1,\sigma^1),( \mu^2,\sigma^2 )$ are concatenated to form a comparison vector $h$. An attention mechanism operates within each section $i$ to interpolate between the $M$ available generators, producing a composite direction $\mathfrak{g}_c^i$. Sectionwise on/off selection is implemented with Gumbel-Softmax switches $sw^i$, informed by pseudo-labels derived from latent coordinate differences, so that only factors undergoing change are activated. The resulting group element is

$$g_c = \exp\Bigl(\sum_{i=1}^K sw^i \mathfrak{g}_c^i\Bigr)$$

and maps $z^1$ to an estimate of $z^2$. This construction enables compositionality, equivariance, and data-driven selection of which factors contribute to a given sample transformation, without explicit knowledge of ground-truth factors.

4. Encoder and Decoder Equivariance Integration

The CVAE architecture is structurally conventional, employing standard convolutional encoders and decoders (or optionally, a Spatial-Broadcast decoder). Uniquely, group equivariance conditions are imposed: $q_\phi(\psi_c \ast x^1) \approx g_c q_\phi(x^1)$

$$p_\theta(g_c z^1) \approx \psi_c \ast p_\theta(z^1)$$

where $\psi_c$ is the (unknown) input-space symmetry corresponding to $g_c$. These constraints are enforced softly via mean-square penalties ($\mathcal{L}_{ee}$ on the latent and $\mathcal{L}_{de}$ on the decoded reconstructions). Data pairs are supplied to the same encoder and decoder parameterization, eschewing the need for auxiliary towers or siamese computation streams. This design induces equivariance between input and latent/decoded transformations.

5. Algebraic Regularization and Disentanglement Losses

The symmetry codebook and factor organization are enforced with several loss terms:

• Commutativity ($\mathcal{L}_c$): All learned Lie algebra elements are constrained to commute, ensuring group compositionality: $\exp(A)\exp(B)=\exp(A+B)$.
• Parallelism within factors ($\mathcal{L}_{pl}$): Generators within the same factor section are driven to be parallel.
• Orthogonality across factors ($\mathcal{L}_{pd}$): Sections corresponding to different factors are incentivized to be orthogonal.
• Sparsity / Axis-Alignment ($\mathcal{L}_s$): Each one-parameter subgroup is regularized to move only one latent coordinate, establishing axis-alignment.
• Factor prediction ($\mathcal{L}_{p}$): The section on/off mechanism is trained with cross-entropy against pseudo-labels from latent means.
• Encoder and Decoder Equivariance ($\mathcal{L}_{ee}, \mathcal{L}_{de}$): Penalize deviations from algebraic group action equivariance in both encoding and decoding.

Collectively, these losses effect a disentangled, group-structured latent representation aligned with independent generative axes, without recourse to supervision or labeled factor information.

6. Training Procedure

The CFASL training protocol for the CVAE is as follows:

1. Sample a minibatch of data ${x_i}_{i=1}^N$.
2. Pair samples into $(x^1, x^2)$.
3. For each pair:
   • Encode $$(\mu^1, \sigma^1) = q_\phi(x^1),\ (\mu^2, \sigma^2) = q_\phi(x^2)$$.
   • Draw latent codes $z^1, z^2$ from respective posteriors.
   • Compute attention, switches, and aggregate composite symmetry.
   • Decode reconstructions and apply symmetry transformation in latent.
4. Compute all terms in the augmented loss as above.
5. Update all model and codebook parameters by backpropagation.
6. Iterate until convergence.

The method is designed to induce symmetry-based disentanglement in fully unsupervised fashion, without ever specifying factor names or identities during training (Jung et al., 2024).

7. Multi-Factor Disentanglement Evaluation: m-FVM$_k$

To quantify multi-factor disentanglement, an extended evaluation metric, multi-factor fixed-variance match (m-FVM$_k$), generalizes the Factor-VAE metric to scenarios where $k>1$ factors are held fixed. For each sub-experiment:

• $k$ factors are held constant.
• Minibatches are constructed so that these factors do not vary, the standard deviation of each latent coordinate is computed, and the $k$ smallest-variance coordinates are compared to the known fixed factors.
• The score tallies coincidences with ground-truth over all factor combinations and training epochs.

This metric rigorously measures disentanglement performance when controlling for multiple simultaneous generative axes and confirms that the symmetry-conditioned VAE recovers latent disentanglement under both single- and multi-factor variations (Jung et al., 2024).