SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We advocate for this terminology since previous names such as "mutual information gap" [11] and "mutual information ratio" [72] are ambiguous, and indeed the former of these works considered solely compactness and the latter solely modularity, with neither mentioning the distinction. We remark that when  $n_z > n_s$  (after pruning inactive latents), it is impossible to achieve both perfect modularity and perfect compactness. Of the two, modularity should be prioritized [58, 10] and indeed has been referred to as disentanglement itself [17].
4.2 Explicitness
Modularity and compactness are measured in terms of mutual information, so they are agnostic to how this information is encoded. Our third metric, explicitness, measures the extent to which the relationship between the sources and latents is simple (e.g., linear [41, 17, 59, 18]). Since previous explicitness metrics have been rather ad hoc, we propose a formalism using the framework of predictive  $\mathcal{V}$ -information, a generalization of mutual information that specifies an allowable function class, denoted  $\mathcal{V}$ , for the computation of information [75]. We first estimate the predictive  $\mathcal{V}$ -information of each source  $\mathbf{s}_i$  given all latents  $\mathbf{z}$ :
$$
I _ {\mathcal {V}} (\mathbf {z} \rightarrow \mathbf {s} _ {i}) := H _ {\mathcal {V}} \left(\mathbf {s} _ {i} \mid \varnothing\right) - H _ {\mathcal {V}} \left(\mathbf {s} _ {i} \mid \mathbf {z}\right). \tag {10}
$$
This requires estimating the predictive conditional  $\nu$ -entropy
$$
H _ {\mathcal {V}} \left(\mathbf {s} _ {i} \mid \mathbf {z}\right) := \inf  _ {f \in \mathcal {V}} \mathbb {E} _ {s \sim p (\mathbf {s}), z \sim p (\mathbf {z} \mid \mathbf {s})} [ - \log p (s _ {i} \mid f (z)) ] \tag {11}
$$
and the marginal  $\nu$ -entropy of the source
$$
H _ {\mathcal {V}} \left(\mathbf {s} _ {i} \mid \varnothing\right) := \inf  _ {f \in \mathcal {V}} \mathbb {E} _ {s \sim p (\mathbf {s})} \left[ - \log p \left(s _ {i} \mid f (\varnothing)\right) \right], \tag {12}
$$
where  $\varnothing$  is an uninformative constant. The predictive conditional  $\nu$ -entropy measures how well a source,  $\mathbf{s}_i$ , can be predicted by mapping the latents,  $\mathbf{z}$ , through a function in function class  $\nu$ . Note that this estimation uses the best in-sample negative log likelihood. We choose  $\nu$  to be the space of linear models (though one could pick  $\nu$  to fit particular needs) and so use logistic regression (linear regression) for discrete (continuous) sources. We use no regularization. We compute the marginal  $\nu$ -entropy  $H_{\nu}(\mathbf{s}_i \mid \varnothing)$  in the same way, but substituting a universal constant for all inputs. We propose a simple normalization analogous to the one done for NMI:
$$
\operatorname {N M I} _ {\mathcal {V}} \left(\mathbf {z} \rightarrow \mathbf {s} _ {i}\right) := \frac {I _ {\mathcal {V}} \left(\mathbf {z} \rightarrow \mathbf {s} _ {i}\right)}{H _ {\mathcal {V}} \left(\mathbf {s} _ {i} \mid \varnothing\right)}, \tag {13}
$$
which can be interpreted as the relative reduction in the  $\nu$ -entropy of a source achieved by knowing the latents, and is in [0, 1]: for classification negative log likelihood is the cross-entropy, and for regression we leverage Propositions 1.3 and 1.5 from Xu et al. [75] to argue that  $\mathrm{NMI}_{\nu}(\mathbf{z}\rightarrow \mathbf{s}_i) = R^2$ , the coefficient of determination. We can now compute explicitness as:
$$
\operatorname {I n f o E} := \frac {1}{n _ {s}} \sum_ {i = 1} ^ {n _ {s}} \mathrm {N M I} _ {\gamma} (\mathbf {z} \rightarrow \mathbf {s} _ {i}). \tag {14}
$$
4.3 Summary and Comparison to Nonlinear DCI
We have derived three metrics for evaluating the modularity, explicitness, and compactness of a representation. Each metric has a straightforward information-theoretic interpretation and all share a range of  $[0,1]$ . We order them in decreasing importance and collectively refer to them as InfoMEC := (InfoM, InfoE, InfoC).
Nonlinear DCI [17], a widely used three-pronged framework that measures similar disentanglement properties, suffers several practical drawbacks in comparison to InfoMEC from being defined in terms of relative counts of decision tree splits: determining this requires considering all latents jointly while fitting  $p(\mathbf{s}_i \mid \mathbf{z})$ . This results in a cumbersome computational footprint that is exacerbated by highly sensitive hyperparameters such as tree depth [17, 10], the tuning of which has even seen omission in prior work [47]. These drawbacks worsen with increased latent space dimensionality. In contrast, InfoMEC avoids these issues as it isolates InfoM and InfoC from the choice of predictive function class and only computes pairwise interactions between individual sources and latents. (InfoE also fits  $p(\mathbf{s}_i \mid \mathbf{z})$ , but does so with function classes of severely limited capacity for which fitting procedures scale well.) See Appendix B for experimental vignettes demonstrating the hyperparameter sensitivity of nonlinear DCI (Figure 7) and the robustness of InfoMEC (Figure 6).
Experimental design. We design our experiments to answer the following questions: Does latent quantization improve disentanglement? How does it compare against the strongest known methods that operate under the same assumptions? And, finally, which of our design choices were critical? We benchmark on four established datasets: Shapes3D [9], MPI3D [20], Falcor3D [53], and Isaac3D [53]. Each consists of RGB image observations generated (near-)noiselessly from categorical or discretized numerical sources. Shapes3D is toyish, but the others are chosen for their difficulty [54, 20]. In particular, we use the complex shapes variant of MPI3D collected on a real world robotics apparatus. See Appendix C.1 for further dataset details. Aside from baseline AE and InfoGAN (specifically InfoWGAN-GP, a Wasserstein GAN [1] with gradient penalty [23]), we compare to  $\beta$ -VAE [25],  $\beta$ -TCVAE [11], and BioAE [72], the strongest methods from prior work that obey our problem assumptions and make design decisions mutually exclusive with latent quantization. We also compare to VQ-VAE with  $d = 64$  and  $n_v = 512$  (Figure 2). We ablate weight decay, scalar codebooks, and dimension-specific codebooks from QLAE and weight decay from QLInfoWGAN-GP. We quantify modularity, explicitness, and compactness using both InfoMEC and nonlinear DCI. We qualitatively inspect representations via decoded latent interventions (Figure 4 and Appendix D).
Select experimental details. The choice of decoder architecture is known to impose inductive biases relevant for disentanglement [33, 43]. We use an expressive architecture (Appendix C.3) based on StyleGAN [33, 54] for all methods and datasets. We downsample the observations to  $64 \times 64$  (if necessary). We follow prior work [47] in considering a statistical learning problem rather than a machine learning one: we train on the entire dataset then evaluate on 10,000 i.i.d. samples. We fix
the number of latents in all methods to twice the number of sources. For quantized-latent models, we fix  $n_v = 10$  discrete values per codebook. We tune one key regularization hyperparameter per method per dataset with a thorough sweep (Table 11, Appendix C.2). We use the best performing configurations over 2 seeds and rerun with 5 more seeds. Despite our modest list of methods and ablations, just the last stage took over 1000 GPU-hours.
Table 1: Main disentanglement results measured in InfoMEC and nonlinear DCI. Modularity is the key property, followed by explicitness, with compactness (grayed) a distant third. AE and InfoGAN variants are presented and bolded separately as all AEs are filtered for near-perfect data reconstruction, whereas InfoGANs are generally more lossy. For confidence intervals and data reconstruction results, see Appendix E.
model	aggregated	Shapes3D	MPI3D	Falcor3D	Isaac3D
InfoMEC := (InfoM InfoE InfoC) ↑
AE	(0.40	0.81	0.26)	(0.41	0.98	0.28)	(0.37	0.72	0.36)	(0.39	0.74	0.20)	(0.42	0.80	0.21)
β-VAE	(0.59	0.81	0.55)	(0.59	0.99	0.49)	(0.45	0.71	0.51)	(0.71	0.73	0.70)	(0.60	0.80	0.51)
β-TCVAE	(0.58	0.72	0.59)	(0.61	0.82	0.62)	(0.51	0.60	0.57)	(0.66	0.74	0.71)	(0.54	0.70	0.46)
BioAE	(0.54	0.75	0.36)	(0.56	0.98	0.44)	(0.45	0.66	0.36)	(0.54	0.73	0.31)	(0.63	0.65	0.33)
VQ-VAE	(0.58	0.81	0.39)	(0.72	0.97	0.47)	(0.43	0.57	0.22)	(0.61	0.83	0.42)	(0.57	0.87	0.45)
QLAE (ours)	(0.76	0.84	0.50)	(0.95	0.99	0.55)	(0.61	0.63	0.51)	(0.71	0.77	0.44)	(0.78	0.97	0.49)
InfoWGAN-GP	(0.50	0.57	0.29)	(0.61	0.78	0.41)	(0.43	0.40	0.20)	(0.44	0.60	0.30)	(0.53	0.51	0.24)
QLInfoWGAN-GP (ours)	(0.63	0.59	0.47)	(0.73	0.75	0.48)	(0.62	0.51	0.37)	(0.54	0.53	0.56)	(0.63	0.58	0.49)
DCI := (D I C) ↑
AE	(0.12	0.81	0.10)	(0.11	0.82	0.08)	(0.15	0.83	0.14)	(0.08	0.76	0.07)	(0.13	0.85	0.10)
β-VAE	(0.37	0.89	0.30)	(0.61	0.99	0.47)	(0.31	0.83	0.27)	(0.32	0.84	0.28)	(0.23	0.88	0.19)
β-TCVAE	(0.31	0.87	0.27)	(0.46	0.99	0.38)	(0.22	0.77	0.21)	(0.36	0.90	0.33)	(0.19	0.84	0.16)
BioAE	(0.29	0.86	0.23)	(0.33	0.93	0.25)	(0.24	0.79	0.19)	(0.21	0.81	0.17)	(0.38	0.91	0.31)
VQ-VAE	(0.28	0.79	0.27)	(0.40	0.84	0.34)	(0.09	0.63	0.14)	(0.30	0.79	0.29)	(0.33	0.89	0.31)
QLAE (ours)	(0.59	0.95	0.47)	(0.81	0.99	0.61)	(0.36	0.85	0.36)	(0.50	0.96	0.38)	(0.69	0.99	0.54)
InfoWGAN-GP	(0.14	0.72	0.12)	(0.23	0.80	0.18)	(0.09	0.63	0.09)	(0.11	0.74	0.08)	(0.13	0.71	0.11)
QLInfoWGAN-GP (ours)	(0.26	0.77	0.26)	(0.38	0.85	0.29)	(0.24	0.71	0.25)	(0.20	0.73	0.24)	(0.24	0.79	0.25)
Effect of latent quantization. Adding latent quantization to AE and InfoWGAN-GP results in consistent and dramatic increases in modularity and compactness under both InfoMEC and DCI evaluation (Table 1). For explicitness, the improvement is more significant under (random forest) I than under (linear) InfoE. Qualitatively, decoded latent interventions demonstrate that the QLAE latent space is highly interpretable, corroborating the gains in modularity (Figure 4 and Appendix D).
Comparison of QLAE with prior methods. QLAE significantly outperforms all prior methods on all four datasets in modularity under both InfoM and D (Table 1), with the exception of  $\beta$ -VAE and  $\beta$ -TCVAE on Falcor3D under InfoM. There is also significant improvement in explicitness under (random forest) I, but less so for (linear) InfoE. The objectives in  $\beta$ -VAE and  $\beta$ -TCVAE contain a term that explicitly minimizes the total correlation (aka multiinformation), amongst the latent variables [11]. Since this essentially optimizes for compactness, we should expect compactness metrics to rank  $\beta$ -VAE and  $\beta$ -TCVAE ahead of QLAE; InfoC does so, but C does not. Corresponding metrics from InfoMEC and DCI have Spearman rank correlations of  $\rho = 0.80, p < 1 \times 10^{-11}$  for modularity,  $\rho = 0.77, p < 1 \times 10^{-9}$  for explicitness, and  $\rho = 0.59, p < 1 \times 10^{-5}$  for compactness.
Table 2: Ablation studies on QLAE and QLIinfoWGAN-GP.
model	aggregated	Shapes3D	MPI3D	Falcor3D	Isaac3D
InfoMEC := (InfoM InfoE InfoC) ↑
QLAE (ours)	(0.76 0.84 0.50)	(0.95 0.99 0.55)	(0.61 0.63 0.51)	(0.71 0.77 0.44)	(0.78 0.97 0.49)
QLAE w/ global codebook	(0.68 0.80 0.44)	(0.96 0.99 0.48)	(0.54 0.62 0.45)	(0.59 0.74 0.37)	(0.65 0.82 0.46)
QLAE w/o weight decay	(0.63 0.81 0.48)	(0.69 0.99 0.51)	(0.51 0.61 0.48)	(0.65 0.76 0.43)	(0.66 0.89 0.51)
VQ-VAE w/ weight decay	(0.69 0.87 0.43)	(0.80 0.99 0.46)	(0.50 0.81 0.41)	(0.74 0.86 0.40)	(0.73 0.81 0.44)
QLInfoWGAN-GP (ours)	(0.63 0.59 0.47)	(0.73 0.75 0.48)	(0.62 0.51 0.37)	(0.54 0.53 0.56)	(0.63 0.58 0.49)
QLInfoWGAN-GP w/o w.d.	(0.61 0.55 0.44)	(0.66 0.57 0.40)	(0.63 0.55 0.40)	(0.56 0.56 0.52)	(0.59 0.53 0.45)
DCI := (D I C) ↑
QLAE (ours)	(0.59 0.95 0.47)	(0.81 0.99 0.61)	(0.36 0.85 0.36)	(0.50 0.96 0.38)	(0.69 0.99 0.54)
QLAE w/ global codebook	(0.52 0.93 0.41)	(0.83 0.99 0.60)	(0.36 0.86 0.34)	(0.36 0.90 0.26)	(0.53 0.99 0.43)
QLAE w/o weight decay	(0.49 0.94 0.40)	(0.63 0.99 0.47)	(0.30 0.84 0.29)	(0.46 0.97 0.36)	(0.58 0.99 0.48)
VQ-VAE w/ weight decay	(0.43 0.84 0.37)	(0.74 0.99 0.57)	(0.22 0.68 0.20)	(0.41 0.85 0.32)	(0.34 0.85 0.39)
QLInfoWGAN-GP (ours)	(0.26 0.77 0.26)	(0.38 0.85 0.29)	(0.24 0.71 0.25)	(0.20 0.73 0.24)	(0.24 0.79 0.25)
QLInfoWGAN-GP w/o w.d.	(0.19 0.73 0.19)	(0.16 0.71 0.13)	(0.28 0.74 0.23)	(0.14 0.72 0.17)	(0.20 0.77 0.23)
Ablations on latent space design and regularization. We observe that ablating dimension-specific codebooks, weight decay, and scalar codebooks from QLAE each causes a significant drop in InfoM and D (Table 2), verifying the importance of these design decisions. The effect of ablating weight decay from QLinfoWGAN-GP is less pronounced. We note that while VQ-VAE w/ weight decay performs somewhat closely to QLAE in terms of InfoMEC, this is only because we use the categorical codes (as opposed to the high-dimensional vector representation) for evaluation. In addition, the scalar codebook design of QLAE enables meaningful interpolation between discrete values, whereas this is not supported by vector quantization.
QLAE
BioAE
Figure 4: Decoded Isaac3D latent interventions for QLAE and BioAE, the prior method with highest InfoM. In each image block, a single data sample is encoded into its latent representation. In the  $j$ -th column, the  $j$ -th latent is intervened on with a linear interpolation of its range across the dataset. The resulting representation is then decoded. For reference, the true sources for this dataset are object shape, robot  $x$ , robot  $y$ , camera height, object scale, lighting intensity, lighting direction, object color, and wall color. QLAE's representation is highly interpretable: going down each column corresponds to one source (sometimes two) changing in a consistent manner, with all other sources remaining unchanged. In contrast, the BioAE's representation varies more sporadically, and its decoder often generates low-quality samples from post-intervention latents. For qualitative results for multiple data samples from all datasets alongside NMI heatmaps, see Appendix D.
Nonlinear ICA and disentangled representation learning. There is a long history of trying to build interpretable representations that separately represent the sources of variation in a dataset. This goes back to classic work on (linear) ICA [13, 29], and has been known in deep learning as disentanglement [4]. Without further assumptions, nonlinear ICA and its relatives in disentangled representation learning are provably underspecified [30, 47]. Approaches to resolve this indeterminacy include labeling a small number of datapoints [54] or showing pairs of datapoints in which only one or a few sources differ [63, 48]. More in line with our work are approaches that assume additional structure in the data generative process and disentangle by ensuring the representation reflects this structure [74]. Assumptions and methods include: factorized priors [25, 11, 36, 61], biologically inspired activity constraints [72], sparse source variation over time [65, 38], structurally sparse source to pixel influence [57, 31, 78, 52, 8], geometric assumptions on the source to image mapping [64, 77, 26, 22], sparse underlying causal graphs between sources [42], piecewise linearity [37], and hierarchical generation [43]. Many of these ideas are in principle compatible with latent quantization, and we leave discovery of fruitful combinations to future work.
Factorized latent spaces. VAEs that specify an isotropic Gaussian latent prior regularize the marginalized variational distribution (aka aggregate posterior) towards being a factorized distribution [11]. Roth et al. [61] bias latents to have pairwise factorized support via a Hausdorff set distance regularization. For linear ICA, Whittington et al. [72] prove that regularizing latents to be nonnegative and energy-minimizing results in them having factorized support. Differently from all of these works, latent quantization imbues a model with factorized structure by construction, instead of relying on the optimization of regularized objectives to manifest this structure. The favorable disentanglement that QLAE yields over  $\beta$ -TCVAE [11] and BioAE [72] suggests that this strategy is more effective.
Discrete representation learning. Oord et al. [55] first demonstrated the feasibility of discrete neural representation learning at scale, and their techniques have since been broadly applied, e.g., to videos [70, 76], audio [2, 15, 68, 6], and anomaly detection [50]. The following works design discrete representations similarly to how we do, though for purposes other than unsupervised disentanglement. Several works use one scalar codebook per latent dimension to achieve high efficiency in retrieval [3, 67, 32, 73]. Kobayashi et al. [39] disentangle normal and abnormal features in medical images into separate vector codebooks via pixel-space supervision. Liu et al. [46] and Träuble et al. [69] use multiple codebooks with separately parameterized key and value vectors and investigate the effect of discretization in systematic generalization and continual learning, respectively.
We have proposed to use latent quantization and model regularization to impose an inductive bias towards disentanglement that enables our models to outperform strong prior methods. Ablations verify that our main design decisions are critical. We have also synthesized previously proposed ideas for evaluation into InfoMEC, three information-theoretic disentanglement metrics that rectify or sidestep key drawbacks in existing approaches.
While our results are promising, one concern might be that we have overfit our inductive bias to existing disentanglement benchmarks, in which, just like our model, the sources are discrete and the generative process is (near-)noiseless. Our experiments have already demonstrated the ability of latent quantization to represent sources that have more values (up to 40) than the per-dimension codebook size (fixed to 10) via allocating multiple latent dimensions. Future work should strive to construct disentanglement benchmarks that better reflect realistic conditions, e.g. continuous sources.
Beyond the intuitions and connections to related works we have presented, we do not understand why our method performs as well as it does. It may be fruitful to tackle this empirically, e.g. by probing how QLAE distributes data around its latent space, and how weight decay changes this. Achieving satisfactory understanding would enable the field to better position latent quantization within the ongoing body of work that aims to develop generalizable conditions for successful disentanglement.
Lastly, we hope this method, its future versions, and other methods the field develops are able to deliver on the original motivation for disentangled representation learning—to learn human-interpretable representations in complex, real-world situations, and to leverage the interpretability to empower human decision-making. This will require methods that can disentangle out-of-distribution data samples, that work for generic data types, and that can learn compositionally from sparse interactions with data. We suspect that latent quantization may have a role to play in these directions.
We gratefully acknowledge the developers of open-source software packages that facilitated this research: NumPy [24], JAX [7], Equinox [35], matplotlib [28], seaborn [71], and scikit-learn [56]. We also thank Evan Liu, Kaylee Burns, Karsten Roth, Anirudh Goyal, and Cian Eastwood for feedback on previous drafts.
KH was funded by a Sequoia Capital Stanford Graduate Fellowship. Part of WD's work on this project happened during a visit to Stanford funded by the Bogue Fellowship. JCRW was funded by a Henry Wellcome Post-doctoral Fellowship (222817/Z/21/Z). This work was also in part supported by the Stanford Institute for Human-Centered AI (HAI), NSF RI #2211258, Air Force Office of Scientific Research (AFOSR) YIP FA9550-23-1-0127, and ONR MURI N00014-22-1-2740.
[1] Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning, 2017.
[2] Alexei Baevski, Steffen Schneider, and Michael Auli. vq-wav2vec: self-supervised learning of discrete speech representations. In International Conference on Learning Representations, 2020.
[3] Johannes Ballé, Valero Laparra, and Eero P. Simoncelli. End-to-end optimized image compression. In International Conference on Learning Representations, 2017.
[4] Yoshua Bengio. Deep learning of representations: looking forward. In Statistical Language and Speech Processing, 2013.
[5] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432, 2013.
[6] Zalán Borsos, Raphaël Marinier, Damien Vincent, Eugene Kharitonov, Olivier Pietquin, Matt Sharifi, Olivier Teboul, David Grangier, Marco Tagliasacchi, and Neil Zeghidour. AudioLM: a language modeling approach to audio generation. arXiv preprint arXiv:2209.03143, 2022.
[7] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax.
[8] Simon Buchholz, Michel Besserve, and Bernhard Scholkopf. Function classes for identifiable nonlinear independent component analysis. In Advances in Neural Information Processing Systems, 2022.
[9] Chris Burgess and Hyunjik Kim. 3D shapes dataset, 2018. URL https://github.com/deepmind/3dshapes-dataset/.
[10] Marc-André Carbonneau, Julian Zaidi, Jonathan Boilard, and Ghyslain Gagnon. Measuring disentanglement: a review of metrics. IEEE Transactions on Neural Networks and Learning Systems, 2022.
[11] Ricky TQ Chen, Xuechen Li, Roger B Grosse, and David K Duvenaud. Isolating sources of disentanglement in variational autoencoders. In Advances in Neural Information Processing Systems, 2018.
[12] Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, and Pieter Abbeel. Infogan: interpretable representation learning by information maximizing generative adversarial nets. In Advances in Neural Information Processing Systems, 2016.
[13] Pierre Comon. Independent component analysis, a new concept? Signal processing, 36(3): 287-314, 1994.