PartitionVAE: Interpretable Latent Factorization

• PartitionVAE is a deep generative model that partitions latent space into disjoint groups to enforce feature disentanglement and improve interpretability.
• It uses separate encoder pathways and no cross-partition interactions to create distinct latent subspaces for semantic content and modality-specific styles.
• Experimental results on datasets like MNIST show that partitioning leads to robust reconstructions and clear semantic changes, despite trade-offs such as potential blurring from subresolution techniques.

PartitionVAE refers to a class of variational autoencoder models that impose an explicit partitioning of the latent space into disjoint groups of variables, with the central objective of enhancing interpretability and potentially disentangling meaningful factors of variation. Architectural, probabilistic, and experimental variants of PartitionVAE have been developed across unimodal and multimodal settings, with a prominent focus on interpretability and efficient representation (Sheriff et al., 2023, Hsu et al., 2018).

1. Motivation and Theoretical Foundations

Standard VAEs optimize an evidence lower bound (ELBO) over data $x$, introducing a continuous latent vector $z\in\mathbb{R}^d$ with a Gaussian prior $p(z)=\mathcal{N}(0,I)$. While this induces some regularization, individual latent coordinates have no enforced semantic structure or grouping, impeding human interpretability (Sheriff et al., 2023). Empirically, VAEs arrange similar data points near each other in latent space, but meaningful features are often distributed or entangled across multiple coordinates.

PartitionVAE addresses this problem by decomposing the latent space into $P$ non-overlapping partitions, each designed to represent a distinct, potentially interpretable feature group. The model enforces no cross-partition interactions in both the recognition and generative networks, encouraging each partition to specialize. Formally, for latent dimension $d = \sum_{i=1}^P d_i$, the prior and variational posterior both factorize as

$$p(z) = \prod_{i=1}^P p(z^{(i)}), \quad q(z|x) = \prod_{i=1}^P q(z^{(i)}|x),$$

where $z^{(i)} \in \mathbb{R}^{d_i}$.

This grouping mechanism is designed to align contiguous blocks of latent space with joint factors of variation, making scaling or pertubations within a partition correspond to semantically coherent changes in $x$.

2. Model Architecture and Partitioning Strategies

The canonical unimodal PartitionVAE consists of three stages (Sheriff et al., 2023):

• Encoder: Stacks of convolutional (or other) layers extract features and compress to a shared representation $h\in\mathbb{R}^H$.
• Partition Layer: The latent code is split into $P$ partitions. For $z\in\mathbb{R}^d$0 and $z\in\mathbb{R}^d$1 with $z\in\mathbb{R}^d$2 being small neural networks, the sampled latent is

$z\in\mathbb{R}^d$3

and $z\in\mathbb{R}^d$4.

• Decoder: Maps $z\in\mathbb{R}^d$5 through a linear and upsampling stack (e.g., deconvolution layers) to reconstruct the input.

There are no cross-partition dependencies; each $z\in\mathbb{R}^d$6 only depends on $z\in\mathbb{R}^d$7, and the decoder processes the concatenated latent vector as usual.

The multimodal variant (“Partitioned VAE” or PVAE in (Hsu et al., 2018)) further divides latent space into modality-invariant semantic variables $z\in\mathbb{R}^d$8 and modality-specific style codes $z\in\mathbb{R}^d$9, $p(z)=\mathcal{N}(0,I)$0. The generative process posits independent priors and modality-conditional decoders, with factorization: $p(z)=\mathcal{N}(0,I)$1 and parallel encoders for each modality.

3. Objective Function and Probabilistic Formulation

PartitionVAE maximizes the ELBO, with the loss decomposed across all partitions: $p(z)=\mathcal{N}(0,I)$2 Typically, $p(z)=\mathcal{N}(0,I)$3 and the reconstruction loss is weighted higher than KL terms to prioritize fidelity.

In multimodal PVAE (Hsu et al., 2018), the ELBO extends to both modalities and their respective latent style/semantic codes. Additional regularizers enforce coherence between multimodal and unimodal encoders and push semantic codes of different samples apart via contrastive hinges.

The probabilistic independence of partitions (for unimodal) or sem/ style subspaces (for multimodal) directly supports the intended disentanglement.

4. Subresolution and Computational Optimizations

PartitionVAE introduces a “subresolution” mechanism to mitigate training overhead on resource-limited hardware (Sheriff et al., 2023). The decoder is trained to output downsampled images $p(z)=\mathcal{N}(0,I)$4 with lower spatial resolution (e.g., reduction factor $p(z)=\mathcal{N}(0,I)$5), which are then upsampled (typically via bilinear interpolation) for comparison to the full-resolution input. The reconstruction loss is computed after upsampling, backpropagating through the interpolation operation. This strategy reduces the number of parameters and required FLOPs by approximately $p(z)=\mathcal{N}(0,I)$6, at the cost of potentially blurrier output.

5. Experimental Evaluation and Interpretability

PartitionVAE has been empirically evaluated on MNIST (10-dimensional, split into partitions [4,3,3]) and the Sports10 table tennis subset (20-dimensional, [5,5,4,3,2,1]) (Sheriff et al., 2023). Key findings include:

• MNIST: After 25 epochs, PVAE achieves clean reconstructions with mild blur. Traversing individual partitions yields distinct semantic changes: e.g., scaling one partition transforms a ‘9’ into a ‘2’, another partition to a ‘4’, and another to a ‘7’. This demonstrates high interpretability, where each partition correlates with a major digit component.
• Ping-Pong Images: PVAE separates global/static scene elements (contrast, letterboxing) but struggles with dynamic features (players, ball), particularly due to subresolution-induced blurring and limited capacity. Some small partitions have little visible effect.

For the multimodal PVAE (Hsu et al., 2018), the semantic partition $p(z)=\mathcal{N}(0,I)$7 achieves $p(z)=\mathcal{N}(0,I)$899% clustering purity for digit identity, with style partitions capturing modality-specific attributes for both speech and images.

6. Implementation and Hyperparameters

PartitionVAE is typically implemented in PyTorch (Sheriff et al., 2023), with the following hyperparameters:

• Encoder: Conv2d layers, final Linear to $p(z)=\mathcal{N}(0,I)$9.
• Partition Nets: For each partition, $P$0 and $P$1 are small MLPs.
• Decoder: Linear, reshaped conv-tensor, ConvTranspose layers with activations.
• Training: Adam optimizer, learning rate $P$2, batch size 128 (MNIST), 64 (table tennis), $P$3–$P$4 epochs, reconstruction:KL loss ratio $P$5.
• Subresolution: Downsampling factor $P$6 in experiments.

In multimodal PVAE (Hsu et al., 2018), architectures include LSTMs for speech, convnets for images, and flexible encoder/decoder heads for each partition.

7. Limitations, Trade-offs, and Directions

PartitionVAE offers improved interpretability by trading off some expressive power due to forced factorization and no cross-partition communication. Trade-offs include:

• Subresolution: Yields %%%%33$d = \sum_{i=1}^P d_i$34%%%% training speedup, but induces blur and loss of fine details.
• Partition Redundancy: Larger numbers of partitions may be underutilized; small partitions sometimes have negligible effect.
• KL Weighting: Excessive KL weight risks posterior collapse; most effective settings weigh reconstruction loss higher.

Notable limitations include challenges with complex, high-frequency features and the need to hand-select partition sizes. Potential extensions identified include adaptive partition sizing, learned upsampling schemes, and integrating contrastive/adversarial objectives to safeguard high-frequency detail (Sheriff et al., 2023).

8. Relation to Other Partition-Based VAEs

Partitioned latent space approaches extend beyond PVAE. For example, SepVAE uses partitioning to separate pathological from healthy factors (Louiset et al., 2023), and multimodal PVAE (Hsu et al., 2018) explicitly splits semantics and modality-specific styles, leveraging labeled and unlabeled data to reinforce these separations.

A distinguishing feature of the original PVAE is the explicit “no cross-talk” partition, where internal connectivity is enforced only within partitions. Results suggest that, for simple or highly structured data (e.g., MNIST), this prior yields disentangled, semantic factors, whereas for complex scenes, it primarily captures global attributes.

---

PartitionVAE thus represents a systematic approach to imposing human-interpretable structure within deep generative models by latent-space partitioning, yielding benefits for interpretability and factor disentanglement, but also highlighting inherent trade-offs between statistical efficiency and manual structure enforcement (Sheriff et al., 2023, Hsu et al., 2018).