Disentanglement via Latent Quantization
Authors: Kyle Hsu
Will Dorrell
James C. R. Whittington
Jiajun Wu
Chelsea Finn
Venue: NeurIPS 2023 poster (2023)
Keywords: disentanglement
unsupervised learning
quantization
TL;DR: Quantizing an autoencoder's latent space with discrete scalar codebooks comprises a strong inductive bias towards disentanglement.

Abstract: In disentangled representation learning, a model is asked to tease apart a dataset's underlying sources of variation and represent them independently of one another. Since the model is provided with no ground truth information about these sources, inductive biases take a paramount role in enabling disentanglement. In this work, we construct an inductive bias towards encoding to and decoding from an organized latent space. Concretely, we do this by (i) quantizing the latent space into discrete code vectors with a separate learnable scalar codebook per dimension and (ii) applying strong model regularization via an unusually high weight decay. Intuitively, the latent space design forces the encoder to combinatorially construct codes from a small number of distinct scalar values, which in turn enables the decoder to assign a consistent meaning to each value. Regularization then serves to drive the model towards this parsimonious strategy. We demonstrate the broad applicability of this approach by adding it to both basic data-reconstructing (vanilla autoencoder) and latent-reconstructing (InfoGAN) generative models. For reliable evaluation, we also propose InfoMEC, a new set of metrics for disentanglement that is cohesively grounded in information theory and fixes well-established shortcomings in previous metrics. Together with regularization, latent quantization dramatically improves the modularity and explicitness of learned representations on a representative suite of benchmark datasets. In particular, our quantized-latent autoencoder (QLAE) consistently outperforms strong methods from prior work in these key disentanglement properties without compromising data reconstruction.

Disentanglement via Latent Quantization
Kyle Hsu
Will Dorrell James C. R. Whittington Jiajun Wu Chelsea Finn
OOffrd University Okylehu,,jaajun cbfinn) stanford. (dorrellwec, jewshittinggonnn@gaalliii com
*ttanford University UUniversity College London
Abstract
In disentangled representation learning, model is asked to tease apart dataset's underlying sources of variation and represent them independently of one another. Since the model provided with ground truth information about these sources, inductive biases take paramount role in enabling disentanglement, In this work, we construct inductive bias towards encoding nized latent space. Concretely, we do this by (i) quantizing the latent space into discrete code vectors with separate learnable scalar codebook per dimension and (ii) applying strong model regularization via an unusually high weight decay. Intuitively, the latent space design forces the encoder to combinatorially construct codes from small number of distinct scalar values, which in turn enables the decoder to assign a consistent meaning each value. Regularization then serves to drive the model towards this parsimonious strategy. We demonstrate the broad applicability of this approach by adding to both basic data-reconstructing (vanilla autoencoder) and latent-reconstructinn (InfoGAN) generative models. For reliable evaluation, we also propose InfoMEC, a new set of metrics for disentanglement that is cohesively grounded in information theory and fixes well-established short- comings previous metrics. Together with regularization, latent quantization dramatically improves the modularity and explicitness of learned representations on a representative suite of benchmark datasets. In particular, our quantized-late autoencoder (QLAE) consistently outperforms strong methods from prior work in these key disentanglement properties without compromising data reconstruction.
and decoding from orga-
82
0g
o"
0
p
v.s.
p
V.
01
??
21
21
(a)
Boompie
(b)
Serepe
(c)
(d)
(e)
Figure 1: Motivating observations and illustration of the inductive bias of latent quantization. The mapping between true sources (c) and data (b) is simpler than most other possible generative functions from unstructured spaces (a). To help recover the sources, we use quantized latent codes (d) continuuus codes (dots) are mapped (via arrows) to their nearest discrete codes (black circles), each of which constructed combinatorially from per-dimension scalar codebooks (turquoise and orange ticks). One way this design's inductive bias manifests is in how the codes change from previous values (grayed) as a result codebook optimization: the combinatorially defined codes move in lockstep, changing the quantized representations of many datapoints. In contrast, naively quantizing into individual vector embeddings (e) would result in codebook optimization having comparatively local effect.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Introduction
Our increasing reliance on black-box methods for processing high- -igh-dimensional data underscores the importance of developing techniques for learning human- interpeieeeeretablle representations, To name but a few possible benefits, such representations could foster more informed human decision- making [62, 27], facilitate efficient model debugging and improvement [16, 44], and streamline auditing and regulation [51, 14]. In this context, disentangled representation learning serves as a worthwhile scaffolding: loosely speaking, its goal is for a model to tease apart dataset's underlying sources of variation and represent them independently of one another.
Accomplishing this, however, has proven difficult. Conceptually, the field lacks a formal problem statement that resolves fundamental ambiguities [30, 47] without overly restrictive assumptions (reviewed in Section 6); methodologically. evaluation metrics have been found to be sensitive to hyperparameters, ad hoc, and/or sample inefficient [10]: and empirically, there remains need for an inductive bias that enables consistently good performance in the purely unsupervised setting.
In this work, answer the call for a better inductive bias for disentanglement. Our solution is moti- vated by observing that many datasets of interest are generated from their sources in a compositional manner, which entails a neatly organized source space (Figure 1). This distinguishing property of realistic generative processes applies to real world physics as well as human approximations thereof (e.g., rendering). Hence, broad strategy uncover the true underlying sources is to bias the model towards encoding to and decoding from a similarly structured latent space.
We manifest this inductive bias by drawing from two common ideas in the machine learning literature: discrete representations [55] and model regularization. Specifically, we propose (i) quantizing a model's latent representation into learnable discrete values with a separate scalar codebook per dimension and (ii) applying strong regularization via an unusually high weight decay [49]. Intuitively, forcing the model to use small number of scalar values combinatorially construct many latent codes encourages it to assign a consistent meaning to each value, an outcome that weight decay explicitly incentivizes by regularizing towards parsimony.
side benefit of models with quantized latents is that they sidestep one issue that has hindered evaluation in previous works: they enable the use of simpler, more robust distribution estimation techniques for discrete variables. As a further methodological contribution, we present InfoMEC, a new set of metrics for the modularity, explicitness, and compactness of (both continuous and discrete) representations that is cohesively grounded in information theory and fixes other well-established shortcomings in existing disentanglement metrics.
We demonstrate the broad applicability of latent quantization by adding it to both basic data- reconstructing (vanilla autoencoder) and latent-reconstructing (InfoGAN) generative models. To- gether with regularization, this is sufficient to dramatically improve the modularity and explicitness of the learned representations of a representative suite of four disentangled representation learning datasets with image observations and ground- ground-truth source evaluations [9, 20, 53]. In particular, our quantized-late latten autoencoder (QLAE, pronounced like clay) consistently outperforms strong methods from prior work without compromising data reconstruction. We think of QLAE as a minimalist implementation of a combinatorial representation in neural networks, suggesting that our recipe of latent quantization and regularization could be broadly useful to other areas of machine learning.
2 Preliminaries
In order to properly contextualize our proposed inductive bias, methodological contributions, and experiments, we first devote some attention to explaining the problem of disentangled representation learning. We also discuss prior disentangled representation learning methods we build upon.
2.1 Nonlinear ICA and Disentangled Representation Learning
We begin by considering the standard data-gen gettarannn model of nonlinear independent components analysis (ICA) [30], a problem very related to but more conceptually precise than disentanglementt
(ss) = II p(s:). x gs)), 1=1
'Code for models and InfoMEC metrics: httpii//htttpp/ggggbb..ooo./ http://pww...ggt......./..aaa/aa

---

where s = (si..... comprise the n, mutually independent source variables (sources); x is the observed data variable; and g : S - X' is the nonlinear data-generating function. The nonlinear ICA problem is to recover the underlying sources given dataset 'D of samples from this model. Specifically, the solution should include an approximate inverting function g - 2 such that, assuming latent variables (latents) z = (Z1..... ) are matched correctly to the sources, each source is perfectly determined by its corresponding latent. A typical technical phrasing is that g- og should be the composition of permutation and dimension-wise invertible function.
As stated, this problem is nonidentifiable (or underspecified). Given D. one may find many sets of independent latents (and associated nonlinear generators) that fit the data despite being non-trivially different from the true generative sources [30]. As such, reliably recovering the true sources from the data is impossible. Much recent work in nonlinear ICA has focused on proposing additional problem assumptions so as to provably pare down the possibilities to unique solution. These theoretical assumptions can then be transcribed into architectural choices or regularization terms. Such approaches have shown promise in increasing our understanding of what assumptions are required to disentangle; we review these works in Section 6.
While identifiability is conceptually appealing, achieving under sufficiently generalizable assump- tions that apply to non-toy datasets has proven hard. The field of disentangled representation learning has taken a more pragmatic approach, focusing on empirically evaluating the recovery of each dataset's 's designated source set. Given this empirical focus, robust performance metrics are vital. Unfortunately, there is a plethora of approaches in use, and subtle yet impactful issues arise even in the most common choices [10]. To address these concerns, in Section 4 we propose new metrics for three existing, complementary notions of disentanglement [17, 59]. They measure the follow- ing three properties: modularity- the extent to which sources are encoded into disjoint sets of latents; explicitness simply the latents encode each source; and compactness the extent to which latents encode information about disjoint sets of sources. We frame these three in cohesive information-theoretie framework that we name InfoMEC.
The disentangled representation learning problem statement considered in this work is as follows. Given a dataset of paired source-date samples {(s, x = g(s))} from the nonlinear ICA model (1), learn an encoder 9 X' + Z and decoder 9 Z + X' solely using the data {z} such that (i) the InfoMEC as estimated from samples {(8, = = 9 o (()))} from the joint source-latent distribution is high, while (ii) maintaining an acceptable level of reconstruction error between x and I 9 (x).
2.2 Autoencoding and InfoGAN as Data and Latent Reconstruction
We will apply our proposed latent quantization scheme to two foundational approaches for disen- tangled representation learning: vanilla autoencoders (AEs) and intommation-mximizze generative adversarial networks (InfoGANs) [12]. Here, we provide a brief overview of the two and defer com- plete implementation details Appendices A and C. Both approaches involve learning an encoder g and decoder g. An autoencoder takes a datapoint x E X as input and produces a reconstruction g o 1--2(() (xr) E X that is optimized match the input:
Preconstrnct data (g- 0 D) :- Eg-D [- log و (x))]
An InfoGAN instead takes a latent code E Z as input. The decoder (aka generator) maps to the data space, and from this the encoder produces a reconstruction of the latent:
Creconstractt latem (g . = Fpppp()) [- logg g g(=))]
Unlike the data reconstruction loss, this is clearly insufficient for learning as the dataset D isn't 't even used. InfoGAN can be thought of grounding latent reconstruction by making the marginal distribution of generated datapoints, (z). indistinguishable from the empirical data distribution. A concrete measure of this is provided by an additional binary classifier (aka discriminator) [21] or value model (aka critic) [1] trained alongside but in opposition to the decoder. While InfoGAN was originally motivated as maximizing a variational lower bound on the mutual information between the latent and the generated data, we find the above interpretation to be unifying.
3 Latent Quantization
Our goal is to encourage our model to disentangle by biasing it towards using an organized latent space. Why would this mitigate the nonidentifabiility of nonlinear ICA? Our key motivation is that
3

---

212
22,2
22
0
21,1
22,1
21
(a) vector quantization
(b) latent quantization
Two ways to quantize into one of n. n= (52) discrete codes. (a) Vector quantization [55] splits continuous representation of size n.d (4) into parts (black dots), each of which is quantized the nearest of global codebook of n, vectors of size d (turquoise circles). (b) Latent quantization specifies d 1 and quantizes a continuous representation of size n. (black dot) onto regular grid parameterized by dimension-specific codebooks (orange and turquoise ticks), each of size ny. Unlike vector quantization, latent quantization ties the decoder input space to the discrete code space.
Figure 2:
generative processes for realistic data are compositional and hence necessarily use highly organized source spaces. We discuss connections to related works in Section 6.
To build our desired inductive bias, propose latent quantization, a modification of vector quan- tization [55]. In the latter, the latent representation of a datapoint is partitioned into component vectors, each of which is quantized the nearest of discrete of vector embeddings (Figure 2a). We modify vector quantization to better enable the model to encode and decode with a consistent interpretation of each component of the discrete code. Concretely, we do this by specifying the set of latent codes to be the Cartesian product of distinct scalar codebooks: Z V1 X ... x VR.. where each codebook Vy is set of n, reals. The complete transformation comprises, first, the encoder network mapping the data to continuous latent space 9-1 X - R (with slight abuse of notation), followed by quantization onto the nearest code (Figure 2b). This nearest neighbor calculation be done elementwise, which is highly efficient. Formally, latent quantization is:
2z arg min و (x)) - Uykl, j= ",
We represent Z by its constituent discrete values, concretely learnable two-dimensional array E R", the j-th row which stores elements of codebook V, an arbitrary order.
lesson from nonlinear ICA is that, given a flexible enough model, data can be mapped to and from latent spaces many convoluted ways. We motivate the use of strong model regularization with the conjecture that, of all the possible mappings from organized latent space to data, the most parsimonious will be the true generative model something close enough this by using high weight decay on both the encoder and decoder networks. Ablation studies (Section 5) show that both quantized latents and weight decay are necessary to disentangle well.
We operationalize
To train a quantized latent model, we use the straight-througe gradient estimator [5] and co-opt the quantization and commitment losses proposed for vector quantization:
Lquattize |StooGrrdient(g (x)) =113.
Lcommit |1 (xx StopGradient(3)13 (5)
The ttraight-through gradient estimator facilitates the flow of gradients through the nondifferentiable quantization step. Lquantiz pulls the discrete values constituting = onesidedly towards the pre- quantized continuous output of the encoder. This needed to optimize V, since straight-through gradient estimation disconnects V from the computation graph. Conversely, the commitment loss prevents the pre-quantized representation, which does see gradients from downstream computation, from straying too far from the codes. While this is significant failure mode for vector quantization, we find that the use of scalars instead of high-dimensional vectors alleviates this issue, allowing us to drastically downweight the quantization and commitment losses while maintaining training stability. This gives the model much-needed flexibility to reorganize the discrete latent space. Finally, while using shared global codebook like in vector quantization is certainly feasible, we find it better to maintain dimension-specific codebooks to enable the stable optimization of each individual value. Algorithm I contains pscudocode for latent quantization and computing the quantization and

---

commitment losses, Appendix A presents pseudocode for training a guantized-latemt autoencoder (QLAB) in Algorithm 2 and a quamtized InfoWGAN-GP [12, 23] in Algorithm 3.
Algorithm Latent quantization and computation of codebook losses.
1: function Lateentuuaniizzaiionndannapo
v; encoder gー discrete value array V)
2:
Zc (x)
3:
2 f- arg min ue mizl vI[1, E J
V V
D implement via (4)
4:
Lquantiz < |StooGradient(2.4 z.) 2
from VQ-VAE [55]
5:
Lcommit
Ec. StopCrradient(2))
D ibid.
6:
z t Ze + StopGradient( ze)
> straight-through gradient estimator [5]
7:
return Qqaatizz , Ccommit
4 InfoMEC: Cntommttion-Thoretii Metrics for Disentanglement
In this section, we derive InfoMEC, metrics for modularity, explicitness, and compactness, building upon and otherwise taking inspiration from several prior works [58, 17, 11, 34, 72, 45, 10]. We take care to motivate our design decisions, and while we do not expect this be the final word on disentanglement metrics, we hope our presentation enables others to clearly understand InfoMEC and propose further improvements.
4.1 Modularity and Compactness
Nonlinear ICA asks for the latents recover the sources up to a permutation and dimension wise invertible transformation. The mutual information between an individual source and latent,
(s:; 23) == (p(s,, 2,) p(2)))()))):
(6)
is granular measure of the extent to which they are deterministic functions of each other. Unlike other measures such as correlation (used in MCC [34]), LASSO weights (used in linear DCI (17)), or linear predictive accuracy (used in SAP [41]), mutual information takes into account arbitrary nonlinear dependence between its two arguments, making it invariant within the nonlinear ICA equivalence class for any candidate solution.
When both arguments are discrete, estimating the mutual information is simple via the empirical joint distribution, but if either is continuous, estimation becomes noivtrivial. Previous works bin continuous variable and pretend is discrete [47, 10), but this sensitive to the binning strategy [10]. Instead, for evaluating continuous latents, we choose the celebrated k-nearest neighbor based KSG estimator [40, 19], particular a variant [60] designed to handle mix of discrete and continuous arguments. We use k = 3. See Appendix B for experimental vignettes demonstrating the severe sensitivity of binning- based estimation to the binning strategy (Figure 5) and the robustness of KSG-based estimation to & (Figure 6). We remark that latent quantization enables reliable evaluation using the discrete-discred estimator.
To facilitate aggregation, we desire a normalization to the interval [o, 1]. To this end, note that the identity [ss 2,) H(,,) H(s, zy) and the nonnegativity of entropy imply [ss; S H(s,) for discrete sources. Following [1]] 1], we define a normalized mutual information as
[(s;: 2z).
NMI(s,, 2,) = H (s:))
(7)
We prefer this normalization scheme over others [66] since i) it is the proportion of source's entropy reduced by conditioning on latent and thus scales consistently to [0, 1] for any model, and ii) avoids the scale-dependent (and possibly negative) differential entropy of continuous latent. We gather all evaluations of NMI(s; Zz)) into a 2-dimensional array NMI e [0, 11m,,,xx We remove inactive latents (columns of NMT), which those with zero range (over the evaluation sample) for discrete latents. For continuous latents, zero too strict, so we heuristically define the threshold to be V20, applied after dividing the ranges by their maximum. See Figure 3 for examples NMIT.
Modularity is the extent to which sources are separated into disjoint sets of latents. Perfect modularity occurs when each latent is informative of only one source, i.e. when every column of NMT has only one nonzero element. This has been measured as the gap [11] between the two largest entries in
5

---

AE
QLAB
8-TCVAR
10
6:10
KCX
6X
ORT
415
25
3
3
2:1
2
RAD
ON
81
R
0:00
OM
5
C:
AST
3
5
X4
16 Figure 3: Visualization of NMT for Shapes3D models. Inactive latents (red font) are removed from InfoM and InfoC computation. The low values in the NMTT of AE indicate that individual latents not highly informative of individual sources. Adding latent quantization results in QLAE achieving near-perfect InfoM (0.99): each latent is only informative of one source (rows are sparse). This induces a trade-off with InfoC, in which 3-TCVAE scores higher (0.53 vs. 0.62). See Appendix D for qualitative studies on the fidelity of NMI entries as judged by decoded latent interventions.
column, or the ratio [72] of the largest entry in the column to the column sum. We prefer the ratio since the gap is agnostic to the smallest nts - 2 values in the column, but these values matter and should influence the measure [10]. Since the possible range of values for this ratio is [//,,, 4], we re-normalize to [0, 1]. Finally, we define InfoModularity (InfoM) as the average of this quantity over latents:
max, NMI the -i=1 NMI.
1 1h
1 n1
InfoM ==
(8)
Compactness complements modularity; it is the extent to which latents only contain information about disjoint sets of sources. We therefore define InfoCompactness (InfoC) analogously to InfoM, but considering rows of NMT instead of columns, and averaging over sources instead of latents, etc.:
max, NMI, j=1 NMI
/ (1-
InfoC: :=
(9)
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. > n, (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, explicitess,, 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 formalism using the framework of predictive V-information, a generalization of mutual information that specifies an allowable function class, denoted V. for the computation of information [75]. We first estimate the predictive V-information of each source s; given all latents z:
[y + s:) := Hy(s: @) Hs z).
(10)
This requires estimating the predictive conditional V-entropy
Hy 2) = inf (e (- logp(s: [(=))
(11)
JEV
and the marginal V-entropy of the source
Hs Ø) = inf E logp(s: f()))),
(12)
jev

---

where Ø is an uninformative constant. The predictive conditional V-entropy measures how well source, s,, can be predicted by mapping the latents, z, through a function in function class V. Note that this estimation uses the best in-sample negative log likelihood. We choose V to be the space of linear models (though one could pick to fit particular needs) and so use logistic regression (linear regression) for discrete (continuous) sources. We use no regularization. We compute the marginal V-entropy fy, Ø) in the same way, but substituting a universal constant for all inputs. We propose simple normalization analogous to the one done for NMI:
NM (z- - s.) lv(z 5s)
(13)
Hy(s; 2)
which can be interpreted as the relative reduction in the V-entropy of a source achieved by knowing the latents, and in [0, 1]: for classification negative log likelihood the cross-entropy, and for regression
we leverage Propositions 1.3 and 1.5 from Xu et al. [75] to argue that NM s;) R²a the coefficient of determination. We can now compute explicitness as:
1 12g
n NMIv(z - 8:).
InfoE:
(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-theoreti 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 (s; z). This results in a cumbersome computational footprint that is exacerbated by highly sensitive hyperparameters such as tree depth [17, 10], the tuming 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 (s; | 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).
5 Experiments
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], MP13D [20], Falcor3D [53], and Isaac3D [53]. from categorical or discretized
Each consists of RGB image observations generated numerical sources. Shapes3D is toyish, but the others are chosen for their difficulty [54, 20]. In par- ticular, we use the complex shapes variant of MPI3D collected on real world robotics apparatus. See Appendix C.1 for further dataset details. Aside from baseline AE and InfoGAN (specifically InfoWGAN-GP, Wasserstein GAN [1] with gradient penalty [23]), we compare to 3-VAE [25], 3-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. 512 (Figure 2). We ablate weight decay, scalar codebooks, and dimension-speciffc codebooks from QLAE and weight decay from JalnfoWGAN-GP.. We quantify modularity, explicitness, and compactness using both InfoMEC and nonlinear DCI. We qualitatively inspect representations via decoded latent interventions (Figure and Appendix D).
Select experimental details. The choice of decoder architecture is known to impose inductive biases relevant for disentanglement [33, 43J. We use an expressive architecture (Appendix C.3) based on StyleGAN [33, 54] for all methods and datasets. We downsample the observations to 64 x 64 (if necessary). We follow prior work [47] in considering a statistical learning problem rather than a machine learning one: train the entire dataset then evaluate 10,000 i..d. samples. We fix
7

---

the number of latents in all methods to twice the number of sources. For quantizee- models, we fix n, = 10 discrete values per codebook. We tune one key regularization hyperparameter per method per dataset with thorough sweep (Table 11, Appendix C.2). We use the best performing configurations over 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 DCL. Modularity 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
Shapes31
MPI3D
Falcor31
Isaac3D
InfoMBC = (InfoM InfoE InfoC) f
AF
(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)
8 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)
B-TCVAR
(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)
BloAE
(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.05 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 (owrs)
(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 (ouns)
(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)
DC1
(DI
(0.12 0.81 0.10)
(0.11
0.82 0.08)
(0.15
0.83 0.14) (0.08
0.76
0.07)
(U.13
0.85
0.10)
G-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)
3-ICVAB
(0.31 0.87 0.27)
(U.46 0.99 0.35)
(0.22 0.77 0.21) (0.36 0.90 0.33)
(U.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-VAR
(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.47)
(0.81
0.99 0.61)
(0.36
0.85 0.36)
(0.50
0.96
0.38)
(0.69
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)
QLinloGGAN 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) InfoF. 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 S-VAE and 3-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 3-E and 3 3-TCVAE contain a term that explicitly minimizes the total correlation (aka multinformationn, amongst the latent variables [1 [1]]. Since this essentially optimizes for compactness, we should expect compactness metrics to rank 3-VAE and 3-TCAA ahead of QLAE; InfoC does so, but C does not. Corresponding metrics from InfoMEC and DCI have Spearman rank correlations of p = 0.80, p < 1 x 10-11 for modularity, P = 0.77, p < 1 x 10 -9 for explicitness, and p = 0.59, p < 1 x 10 for compactness.
Table 2: Ablation studies on QLAE and QL.nfo
model
aggrogated
Shapes3D
MPB3D
Falcor3D
Isac3D
InfoMEC = (JafoM InfoE InfoC) 1
QLAE (ours)
(0.760.8 840 50)
0.995
(0.61 0.63 0.51) (0.71 0.77 0.44) (0.78..97 0. 49)
QLAE global codebook
(0.68 0.80 0.11) 0.9960.99 0848) (0.54 0.62 0.45) (0.39 0.74
(0.65
0.82 0.16)
QLAE w/o weight decay
(0.63 0.81 0 0.48)
(0.69
0.99
(0.51
0.61 0.48) (0.65 0.76 0.43)
(0.66 0.89 0.5)
VQ. VAE w/ weight decay
(0.69 0.87 0.43)
(0.80 0.99 0.46)
(0.50
0.81 0.41)
(0.740.86
00)0
(0.73 0.81 0.44)
QLInfoGAN-GG (ours)
00.630.590
0.47)
(0.730.75
0.48) (0.620.51
0.37) (...440530....
(0.630.58 0.49)
QGANNO wie
(0.61 0.55 0.41) (0.66 0.57 0.40)
(0.630.55 0. 10) (0.560.560...... (3.99..33..3
DCT
(DTC) f
QLAE (ours)
(0.590.95 0.50.44 0.47) 0.199.90.61
(0.36..855 0.36)
(0.90.99900
QLAE w/ glohal codebook
(0.52 0.93 0.41)
00.830.990.600
(0.360.86 0.34) (0.36 0.90 0.26)
0.99 0.43)
QLAE w/o weight decay
(0.49 0.94 0.40)
(0.63 0.99 0.47)
00.30 0. 0.29).99)010 (0.46
0.
(0.58 0.99
VQ-VAE -AF 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 (surs)
(0.260.77 0 .2)
(0.38
0.9 (0. 2. 0. 25) (0. 20..000.73
00.240.790.25
QLInfoWGAN-GE
(0.19 0.73 0.19)
(0.16 0.71 0.13)
0..28 23) (0.140.72 0.17] 20.2000770..2 0.23)
8

---

Ablations on latent space design and regularization. We observe that ablating inntension--seciiie codebooks, weight decay, and scalar codebooks from QLAE each causes significant drop in InfoM and D (Table 2), verifying the importance of these design decisions. The effect of ablating weight decay from QLJnfoWGAN-GP is less pronounced. We note that while VQ- VAE 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. addition, the scalar codebook design of QLAE enables meaningful interpolation between discrete values, whereas this is not supported by vector quantization.
QLAE
Zn
Z4
Zg
2,4
Z15
Zi6
BioAE
Zn
2,
2,
Z,
Z,
2.
2,
2,
24
Zg
Z,0
Z11
Z,2
2,3
Z,4
Z15
Zi6
Z17
Figure 4: Decoded Isaac3D latent interventions for QLAE and BioAE, the prior method with highest InfoM. In each image block, single data sample is encoded into its latent representation. In the j-th column, the j-th latent is intervened with 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 xr 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 one source (sometimes two) changing in consistent manner, with all other sources remaining unchanged. In contrast, the BioAE's '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.

---

6 Related Work
Nonlinear ICA and disentangled representation learning. There is 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, 29I, 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 resolve this indeterminacy include labeling a small number of datapoints [54] or showing pairs of datapoints in which only one or few sources differ 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 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 factorized distribu- tion [11]. Roth al. [61] bias latents have pairwise factorized support via Hausdorff set distance regularization. For linear ICA, Whittington et al. [72] prove that regularizing latents to be nonnegative and energy-mining meergyiiiiiiiziggggg results in them having factorized support. Differently from all of these works, latent quantization imbues 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 3-TCVAE [11] and BioAE [72] suggests that this strategy is more effective.
Discrete representation learning. Oord 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 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.
7 Discussion
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-theorre or sidestep key drawbacks in existing approaches.
disentanglement metries that rectify
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 mnois-ssssssess. 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 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, 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 leverage the interpretability to empower human decision-making. This will require methods that can disentangle out-of-distributic 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 role to play in these directions.
10

---

Acknowledgments and Disclosure of Funding
We gratefully acknowledge the developers of open-source software packages that facilitated this research: NumPy [24]. JAX [7], Equinox [35], matplodli [28], seaborn [71], and scikit-learn [56]. We also thank Evan Liu, Kaylee Burns, Karsten Roth, Anirudh Goyal, and Cian Eastwood for feedback previous drafts.
KH was funded by Sequoia Capital Stanford Graduate Fellowship. Part of WD's work on this project happened during visit to Stanford funded by the Bogue Fellowship. JCRW was funded by Henry Wellcome Post-doctoral Fellowship (22281 12/2/7///21/7)) This work was also part supported by the Stanford Institute for Human-Centered AI (HAD), NSF RI #221 1258, Air Force Office of Scientific Research (AFOSR) YIP FA9550-23-1-00277 and ONR MURI N00014-22-1-2740.
References
[1] Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. International Conference on Machine Learning, 2017.
[2] Alexei Baevski, Steffen Schneider, and Michael Auli. vq-wav2vec: self- self-supervised learning of discrete speech representations. International Conference on Learning Representations, 2020.
[3] Johannes Ballé, Valero Laparra, and Eero P. Simoncelli. End-to-end optimized image compres- sion. 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: 130.. 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. AudioL language modeling approach to audio generation. arXiv preprint arXiv: arXiv:2209.03144 2022.
[7] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Minn Milnee, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
http://githei cwm/ gotggooogllaaaa
[8] Simon Buchholz, Michel Besserve, and Bernhard Schölkopf. 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://gittuub.com deepmind/3dshapen daaassttttt.t....
[10] Marc-André Carbonneau, Julian Zaidi, Jonathan Boilard, and Ghyslain Gagnon. Measuring disentanglement: 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.
11

---

[14] Elliot Creager. David Madras, Joern- Bern--e Jacobsen, Marissa Weis, Kevin Swersky, Toniann Pitassi, and Richard Zemel. Flexibly fair representation learning by disentanglement. In International Conference on Machine Learning, 2019.
[15] Prafulla Dhariwal, Heewoo Jun, Christine Payne, Jong Wook Kim, Alec Radford, and Ilya Sutskever. Jukebox: a generative model for music. arXiv preprint arXiv:2005.00:: 2020.
[16] Mengnan Du, Ninghao Liu, and Xia Hu. Techniques for interpretable machine learning. Communications of the ACM, 63(1):68-77, 2019.
[17] Cian Eastwood and Christopher KI Williams. A framework for the quantitative evaluation of disentangled representations. In International Conference on Learning Representations, 2018.
[18] Cian Eastwood, Andrei Liviu Nicolicioiu, Julius Von Kügelgen, Armin Kekić, Frederik Träuble, Andrea Dittadi, and Bernhard Schölkopf. Dci-es: An extended disentanglement framework with connections to identifiability. In International Conference Learning Representations, 2023.
[19] Weihao Gao, Sewoong Oh, and Pramod Viswanath. Demystifying fixed k-ncarest neighbor information estimators. IEEE Transactions on Information Theory, 64(8):5629-566) 2018.
[20] Muhammad Waleed Gondal, Manuel Wuthrich, Djordje Miladinovic, Francesco Locatello, Martin Breidt, Valentin Volchkov, Joel Akpo, Olivier Bachem, Bernhard Schölkopf, and Stefan Bauer. On the transfer of inductive bias from simulation to the real world: a new disentanglement dataset. In Advances in Neural Information Processing Systems, 2019.
[21] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Fanley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Advances in Neural Information Processing Systems, 2014.
[22] Luigi Gresele, Julius Von Kügelgen, Vincent Stimper, Bernhard Schölkopf, and Michel Besserve. Independent mechanism analysis, a new concept? Advances in Neural Information Processing Systems, 2021.
[23] Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C Courville. Improved training of Wasserstein GANs. In Advances Neural Information Processing Systems, 2017.
[24] Charles R Harris, K Jarrod Millman, Stéfan J Van Der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathanicl J Smith, et al. Array programming with NumPy. Nature, 585(7825):357-362, 2020.
[25] Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. 3-VAE learning basic visual concepts with constrained variational framework. In International Conference on Learning Representations, 2017.
[26] Daniella Horan, Eitan Richardson, and Yair Weiss. When is unsupervised disentanglement possible? In Advances in Neural Information Processing Systems, 2021.
[27] Xiaowei Huang, Daniel Kroening, Wenjie Ruan, James Sharp, Youcheng Sun, Emese Thamo, Min Wu, and Xinping Yi. A survey of safety and trustworthiness of deep neural networks: verification, testing, adversarial attack and defence, and interpretability. Computer Science Review, 37:100270, 2020.
(28) John D Hunter. Matplotlib: 2D graphics environment. Computing in Science & Engineering, 903):90-95, 2007.
[29] Aapo Hyvärinen and Erkki Oja. Independent component analysis: algorithms and applications. Neural Networks, 13(4-5):411 43300 2000.
[30] Aapo Hyvärinen and Petteri Pajunen. Nonlinear independent component analysis: existence
and uniqueness results. Neural Networks, 12(3):429-439 1999.
[31] Li Jing, Jure Zbontar, et al. Implicit rank-minimizing autoencoder. In Advances in Neural Information Processing Systems, 2020.
12

---

[32] Lukasz Kaiser, Samy Bengio, Aurko Roy, Ashish Vaswani, Niki Parmar, Jakob Uszkoreit, and Noam Shazeer. Fast decoding sequence models using discrete latent variables. In International Conference on Machine Learning, 2018
[33] Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Conference on Computer Vision and Pattern Recognition, 2019.
[34] Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational autoen- coders and nonlinear ICA: a unifying framework. In International Conference on Artificial Intelligence and Statistics, 2020.
[35] Patrick Kidger and Cristian Garcia. Equinox: neural networks in JAX via callable PyTrees and filtered transformations. Disserentiable Programming Workshop at Neural Information Processing Systems, 2021.
[36] Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In International Conference on Machine Learning, pages 1849-2658.. PMLR, 2018.
[37] Bohdan Kivva, Goutham Rajendran, Pradeep Ravikumar, and Bryon Aragam. Identifiability of deep generative models without auxiliary information. In Advances in Neural Information Processing Systems, 2022.
[38] David A. Klindt, Lukas Schott, Yash Sharma, Ivan Ustyuzhaminov, Wieland Brendel, Matthias Bethge, and Dylan Paiton. Towards nonlinear disentanglement in natural data with temporal sparse coding. In International Conference Learning Representations, 2021.
[39] Kazuma Kobayashi, Ryuichiro Hataya, Yusuke Kurose, Mototaka Miyake, Masamichi Taka- hashi, Akiko Nakagawa, Tatsuya Harada, and Ryuji Hamamoto. Decomposing normal and abnormal features of medical images for content-based image retrieval of glioma imaging. Medical Image Analysis, 74:102227, 2021.
[40] Alexander Kraskov, Harald Stögbauer, and Peter Grassberger. Estimating mutual information. Physical Review E, 69(6)3066138, 2004.
[41] Abhishek Kumar, Prasanna Sattigeri, and Avinash Balakrishnan. Variational inference of disentangled latent concepts from unlabeled observations. In International Conference on Learning Representationss 2018.
[42] Sebastien Lachapelle, Pau Rodriguez, Yash Sharma, Katie E Everett, Rémi Le Priol, Alexandre Lacoste, and Simon Lacoste-Julien. Disentanglement via mechanism sparsity regularization: new principle for nonlinear ICA. In Conference on Causal Learning and Reasoning, 2022.
[43] Felix Leeb, Giulia Lanzillotta, Yashas Annadani, Michel Besserve, Stefan Bauer, and Bernhard Schölkopf. Structure by architecture: structured representations without regularization. In International Conference on Learning Representations, 2023.
[44] Piyawat Lertvittayakumjorn and Francesca Toni. Explanation- based human debugging of NLP models: survey. Transactions of the Association for Computational Linguistics, 9:1508 -5288, 2021.
[45] Zhiyuan Li, Jaideep Vitthal Murkute, Prashnna Kumar Gyawali, and Linwei Wang. Progressive learning and disentanglement of hierarchical representations. In International Conference on Learning Representations, 2020.
[46] Dianbo Liu, Alex M Lamb, Kenji Kawaguchi, Anirudh Goyal ALIAS PARTH GOYAL, Chen Sun, Michael C Mozer, and Yoshua Bengio. Discrete-valued neural communication. In Advances in Neural Information Processing Systems, 2021.
[47] Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations. In International Conference on Machine Learning, 2019.
[48] Francesco Locatello, Ben Poole, Gunnar Rätsch, Bernhard Schölkopf, Olivier Bachem, and Michael Tschannen. supervssd disentanglement without compromises. In International Conference Machine Learning, 2020.
13

---

[49] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019.
[50] Sergio Naval Marimont and Giacomo Tarroni. Anomaly detection through latent space restora- tion using vector quantized variational autoencoders. In International Symposium on Biomedical Imaging, 2021.
[51] Brent Mittelstadt. Auditing for transparency in content personalization systems. International Journal of Communication, 10:12, 2016.
[52] Gemma Elyse Moran, Dhanya Sridhar, Yixin Wang, and David Blei. Identifialle deep generative models via sparse decoding. Transactions on Machine Learning Research, 2022.
[53] Weili Nie. High resolution disentanglement datasets, 2019. URL https:// http//ggtthttutet com/ WVlabs/High-res haaab///iigg-reeeeddaaaaaaaeeee d
[54] Weili Nie, Tero Karras, Animesh Garg, Shoubhik Debnath, Anjul Patney, Ankit Patel, and Anima Anandkumar. Semi-supervised StyleGAN for disentanglement learning. In International Conference on Machine Learning, 2020.
[55] Aäron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learning. In Advances in Neural Information Processing Systems, 2017.
[56] Fabian Pedregosa, Gaäl Varoquaux, Alexandre Gramfort, Vincent Michel, Bertrand Thirion, Olivier Grisel, Mathieu Blondel, Peter Prettenhofer, Ron Weiss, Vincent Dubourg, Jake Vander- plas, Alexandre Passos, David Coumapeau, Matthieu Brucher, Matthieu Perrot, and Edouard Duchesnay. Scikit-learn: machine learning in Python. Journal of Machine Learning Research, 12(85):2825-2385 2830, 2011.
[57] William Peebles, John Peebles, Jun-Yan Yan Zhu, Alexei Efros, and Antonio Torralba. The Hessian penalty: A weak prior for unsupervised disentanglement. In European Conference Computer Vision, 2020.
[58] Karl Ridgeway, A survey of inductive biases for factorial representatiion-lamiing arXiv preprint arXiv.1612.05299, 2016.
[59] Karl Ridgeway and Michael C Mozer. Learning deep disentangled embeddings with the F-statistic loss. In Advances in Neural Information Processing Systems, 2018
[60] Brian C Ross. Mutual information between discrete and continuous data sets. PLOS ONE, 9(2):
1-5, 2014.
[61] Karsten Roth, Mark Ibrahim, Zeynep Akata, Pascal Vincent, and Diane Bouchacourt. Disentan- glement of correlated factors via hausdorff factorized support. In International Conference on Learning Representations, 2022.
[62] Philipp Schmidt and Felix Biessmann. Quantifying interpretability and trust in machine learning systems. In AAA/ Workshop on Network Interpretability for Deep Learning, 2019.
[63] Rui Shu, Yining Chen, Abhishek Kumar, Stefano Ermon, and Ben Poole. Weakly supervised disentanglement with guarantees. In International Conference on Learning Representations, 2020.
[64] Peter Sorrenson. Carsten Rother, and Ullrich Köthe. Disentanglement by nonlinear ICA with general incompressible-flow networks (GIN). In International Conference on Learning Representations, 2020.
[65] Henning Sprekeler, Tiziano Zito, and Laurenz Wiskott. An extension of slow feature analysis for nonlinear blind source separation. Journal of Machine Learning Research, 15(1):921-947, 2014.
[66] Alexander Strehl and Joydeep Ghosh. Cluster ensembles knowledge reuse framework for combining multiple partitions. Journal of Machine Learning ( 2002.
14

---

[67] Lucas Theis, Wenzhe Shi, Andrew Cunningham, and Ferenc Huszár. Lossy image compression with compressive autoencoders. In International Conference on Learning Representations. 2017.
[68] Andros Tjandra, Sakriani Sakti, and Satoshi Nakamura. Transformer VQ-VAB for unsupervised unit discovery and speech synthesis: ZeroSpeech 2020 challenge. In INTERSPEECH, 2020.
[69] Frederik Träuble, Anirudh Goyal, Nasim Rahaman, Michael Curtis Mozer, Kenji Kawaguchi, Yoshua Bengio, and Bernhard Schölkopf. Discrete key-value bottleneck. In International Conference on Machine Learning, 2023.
[70] Jacob Walker, Ali Razavi, and Aäron van den Oord. Predicting video with VQVAE. arXiv preprint arXiv:2103.019555 2021.
[71] Michael Waskom. Seaborn: statistical data visualization. Journal of Open Source Software, (60):3021, 2021.
[72] James C Whittington, Will Dorrell, Surya Ganguli, and Timothy Behrens. Disentanglement with biological constraints: A theory of functional cell types. In International Conference on Learning Representations, 2023.
[73] Hanwei Wu and Markus Flierl. Learning product codebooks using vector-quantized autoen- coders for image retrieval. In Global Conference on Signal and Information Processing, 2019.
[74] Quanhan Xi and Benjamin Bloem-Reddy, Indeterminacy in generative models: characterization and strong identifiability. In International Conference on Artificial Intelligence and Statisties, 2023.
[75] Yilun Xu, Shengjia Zhao, Jiaming Song, Russell Stewart, and Stefano Ermon. A theory of usable information under computational constraints. In International Conference on Learning Representations, 2020.
[76] Wilson Yan, Yunzhi Zhang, Pieter Abbeel, and Aravind Srinivas. VideoGPT: video generation using VQ- VAE and transformers. arXiv preprint arXiv:2104.10111 2021.
[77] Xiaojiang Yang, Yi Wang, Jiacheng Sun, Xing Zhang, Shifeng Zhang, Zhenguo Li, and Junchi Yan. Nonlinear ICA using volume-preservvng transformations. In International Conference on Learning Representations, 2022.
[78] Yujia Zheng, Ignavier Ng, and Kun Zhang. On the identifiability of nonlinear ICA: sparsity and beyond. In Advances in Neural Information Processing Systems, 2022.
15

---

A
Quantized-Latent Models
This section contains QLnnfOGAAN-G.
pseudocode for latent quantization and for training QLAE and
Algorithm Pseudocode for optimizing a quantized- autoencoder (QLAB).
Hyperparameter values are detailed in Appendix C.2.
Require: dataset D, batch size b, Adam hyperparameters (a, B1, weight decay), loss weights À == (eecoustruut data, Aquaniize Acommit)
1: initialize encoder 9 -1 X' - R=x discrete value array E R",X*., decoder 9 g + X while (g V) has not converged do for b....,ba do
2:
3:
4:
D
5:
2, Lquamiize Leommit f- Latenn(uunnizzaioonn;; 9 V)
D. Algorithm
6:
Creconstruct data BinaryCrossFaatoooy((((.]
(11)
7:
QLAE + X (Creconstmrctt Lquantize, Lcommit)
8:
( (1)) + AdamW (VamW(V(se) (0-1.0) Ei= QLAB: (0 ( -9), B1, 82, weight decay) (W
c(i)
9:
V + Adam( Vv Ei= QLAB a, 31)
Algorithm 3 Pseudocode for optimizing a quantized-laten InfoWGAN-GP.
Hyperparameter values are detailed in Appendix C.2.
Require: dataset D, batch size b, AdamW hyperparameters (a, B1, B2, weight decay), ge--geeeertorrr encoder-to-generatoo
encoder-to-
update ratio The:
loss weights À := (eecmmstruct latent Aquanizze Acommnity Avalues graadient penalty) )
1: initialize encoder g- X' + R=" discrete value array V E R Xn. decoder Rn - X
2: initialize critic 0
R
3: while 0 = (g V, V.9.9. 9-
has not converged do
4: for 1,... Re:g do
b encoder and critic updates
5: for i 1, 1,..., bo
6:
Treal ~ D. Zlake ~ Z, c Uniform((0) 1))
7:
Tlake ((tuke)
8:
Lvalue - gr (Flske) - -.- (
> from WGAN [1]
9:
Tinterpolated CTreal + (1 c)
10:
Egradient penalty &- (I) Eimsspetonn Ir (Finterpolatee) ) |12 1)2 g-1, V)
from WGAN-GP [23]
11:
Ltake; Lequantize Leummit LatentQuantizationnn Preconsruct latent f- Meaan & 7fake)
Algorithm
12:
> from InfoGAN [12]
13:
(() + X (Ereconstmeet latent, Lquantize, Leummit, Livauue Lgradient penalty)
14:
(0 Ir 1)= + AdamW(V (0--0==)) Ei (()) (g 9. 1), a, B1, B2 weight decay)
for 1,..., b
16:
Ztake ~ Z
Ttake t g(Ztaak)
18:
Lvalue &- -0 (Ijake)
>> from WGAN [1]
Etake Eegantiize Leomanit LalennQuaatizzzaiinnniieee 9-1, VV
D Algorithm
20:
Preconstruct latent + MeanSqquaeeEEroor & Z
b from InfoGAN [12]
21:
(1) X (L reecnnstruct latent, Liuantize, Lcommit Lovatue 0)
+ AdamW(V 9 Σ _ 31 B2. weight decay)
23:
V + Adam( (v Et C,,
16

---

B Disentanglement Metries Vignettes
BioAE, S
TCCAE, Falcor3D
AVAE, Isaae3D
InfoM
InfoC
0.4
25
100
50
100
0
50
100
number bins
Figure 5: InfoM and InfoC for three models with continuous latents computed from mutual informa- tion estimated via histogram binning. The binning strategy has drastic effect on the metrics, and is unclear how it should be chosen.
BioAE, S
PTCAAF, Falcor3D
VAE Isaac3D
InfoM
InfoC
25
50
100
25
75
100
25
50
75
100
number of neighbors
Figure 6: InfoM and InfoC for three models with continuous latents computed from mutual informa- tion estimated via l-nniigborrs based KSG [60]. The metrics are very robust the choice of k.
BioAE, Shapes3D
/-TCVAE, Falcor3D
B. VAE, Isaac3D
1.0
0.8
0.6
0.4
0.2
0.0
10
20
0
15
5
15
20
maximum tree depth
Figure 7: Nonlinear DCI three models with continuous latents computed via random forests of 100 trees. D and C are as defined in Eastwood and Williams [17]. I is validation accuracy. DCI is highly sensitive the maximum tree depth hyperparameter, and it easy to overestimate D and this is not tuned with respect to I.
17

---

C Experiment Details
This section contains details the experiments conducted in this work.
C.1
Datasets
Table 3: Summary of datasets used for empirical evaluation.
dataset
n2
[D]
description
Shapes3D
180,000
geometric shapes in
with simplistic textures
MPI3D [20]
460,800
real robot arm holding objects in variety configurations
Falcor3D [53]
7
233,280
single scene illuminated and viewed in varying conditions
Isaac3D [53]
9
737,280
synthetic robot arm holding objects in variety of configurations
Table 4: Shapes3D sources.
index
description
values
0
floor color
10
1
object color
10
2
camera orientation
10
3
object scale
8
4
object shape
5
wall color
15
Figure 8: Shapes3D data samples.
Table 5: MPI3D sources.
index
description
values
0
object color
4
1
object shape
4
object size
3
camera height
3
4 background color
3
5 robot x
40
Figure 9: MPI3D data samples.
6
robot y
40
Table 6: Falcor3D sources.
index
description
values
0
lighting intensity
1
lighting x
6
2
lighting y
6
3
lighting
6
4
camera
5
camera
6
Figure 10: Falcor3D) data samples.
6
camera
6
Table 7: Isaac3D sources.
index
description
values
0
object shape
3
robot x
8
2
robot y
5
3
camera height
4
4
object scale
4
5
lighting intensity
4
lighting direction
6
Figure 11: Isaac3D data samples.
7 object color
4
8 wall color
4
18

---

C.2 Hyperparameters
This section specifies fixed and tuned hyperparameters for all methods considered.
Table 8: Fixed latent quantization hyperparameters for QLAE and QLInfOGGAN-GP.
hyperparameter
value
The VJ initialization
10
linspace( 0.5, 0.5, n.) 1 x 10
quantize
Acomiit
1 10-2
Table Fixed hyperparameters for all autoencoder variants.
hyperparameter
value
f2
2ng
A W learning rate
1 x 10-3
Adam W 31
AddamWa 82
0.99
Adam W updates
< x 105
batch size
128
Areconstruct data
Table 10: Fixed hyperparameters for all InfoGAN variants.
hyperparameter
value
n1 AdamW learning rate
2ng
2 x 10
Adam W 31
AdamW $a
0.9
Adam W updates
< 2 x 105
batch size
64
reconstruct Intent
200
gradient penalty
10
value
Table 11: Key regularization hyperparameter tuning done for each autoencoder and InfoGAN variant.
method
hyperparameter
values
AE
weight decay
0, 0.0.001, 0.01, 0.1]
BioAE [72]
Actiivity
0.01,0.1,1110, 100
8- VAE [25]
AKL
0.1, 0.3, 1,3, 10]
3-TCVAE [11]
8
Arotal correlation
0.3, 1,3, 10
VQ-VAE [55]
weight decay
0.001, 0.001,0.01,0.11/
QLAE (ours)
weight decay
InfoWGAN-GP [54]
weight decay
0.0001100001,00001.000
QLnffWGANN-G (ours)
weight decay
[ 0.0000,00001,00011..11
19

---

C.3 Network Architectures
Inspired by recent works showing how well-designed decoder architectures can facilitate inductive biases relevant for disentanglement [33, 43], we use an expressive architecture for all results presented in this work.
Encoder. We use a simple feedforward convolutional encoder network. Each convolutional block consists of two resolution-preserving convolutional layers (kernll size stride 1) and one downsam- pling convolutional layer (kernel size 4, stride 2) at consistent width (number of channels). Each convolution operation is followed by leaky ReLU (slope 0.3), then instance normalization. There are four such blocks with widths 32, 64, 128, and 256. The 256 x 4 x 4 output is then flattened. Two dense layers each of width 256 with ReLU activation (and normalization) follow. The final operation is an affine projection to the latent layer.
Decoder. We
decoder architecture based StyleGAN [33]. The latent code of shape n; passes through two dense layers each of width 256 with ReLU activation (and no normalization): call the output of this We directly parameterize a starting input feature map of shape 256 x 4 4 with all entries initialized to 0.1. Each style-decoding layer consists of processing an input feature map via a transposed convolution followed by leaky ReLU (slope 0.3) and then adaptive instance normalization (AdalN): w undergoes an affine projection to a scale and bias scalar for each channel of the output feature map, and the output feature map instance normalized then affinely transformed by spatially broadcasting the scales and biases. Similar to the encoder, style-decoding layers are grouped into blocks, with each block consisting of two resolution preseervigg style-decoding layers (kemnel size 3, stride 1) and one upsampling style-decoding layer (kernel size 4, stride 2) at consistent width. There are four such blocks with widths 256, 128, 64, and 32. The final operation is a x 1 convolution of width 3 to yield an output of shape 3 x 64 x 64.
C.4 Negative Results
We tried number of additional modifications our basic methods, QLAE and QLInfoWGAN-GP beyond the ablations presented in the main paper. Here is list of those that only marginally helped, didn't help, made things worse:
Keeping the original VQ-VAE [55] hyperparameter settings for Xquantize (1) and Agommii (0.25).
A step function schedule for the weight decay in AdamW: 0 for the first half of optimization, and a specified value for the second half.
Linearly annealing Aquantiae and/or Acommit from 0 up to a specified value.
Using & norm parameter vector regularization, in addition or without weight decay.
(QLinfWGAN--GP-ppeee Updating V with the encoder instead of with the decoder.
C.5 Nonlinear DCI Evaluation
For each (s; 2) learning problem, we train a random forest classifier with 100 trees and the information gain splitting criterion. We use 0.9/0.1 train/test split of the evaluation sample. We tune the maximum tree depth hyperparameter on held-out accuracy (see Figure for an example of why this is important). We compute D and C using the most predictive model's relative feature importances, following their definitions [17]. We use held-out accuracy for
20

---

D
Qualitative Results
QLAE
B-TCVAE
color
wall color
309
color
6/2
X22
Figure 12: Decoded Shapes3D latent interventions and NMT heatmaps for QLAE and the prior method with highest InfoM. In each image block, single data sample is encoded into its latent representation. In the j-th column, the j-th latent intervened on with linear interpolation of its range across the dataset. The resulting representation is then decoded. NMI heatmaps are annotated with source names, and inactive latents (as determined by range) have red font. In active latents, the visual variation in the generations caused by each latent tightly corresponds sources that have significant NMI values in that latent's 's column. For example, the 3-TCVAE's z has high NMI with object scale and lower NMI with object shape, and going down column in the generations, see consistent changes in object scale and occasional changes in object shape. Inactive latents, e.g. 3-TCVAE's 25, correspond to discernable change in column.
21

---

QLAE
B-TCVAE
objject color
3
object color
030
6:10
000
COK
oiject sloc
AXT
VUST
olject slop
R7A
AXT
XXC
AAT
BXT
beight
height
bacignound color
backgnouund
solus
6M
suboxt
0:00
MP
012
sobory
Figure 13: Decoded MPI3D latent interventions and NMI heatmaps for QLAE and the prior method with highest InfoM. In each image block, single data sample is encoded into its latent representation. In the j-th column, the j-th latent intervened on with a linear interpolation the dataset. The resulting representation is then decoded. NMT heatmaps names, and inactive latents (as determined by range) have red font. In active latents, the visual variation in the generations caused by each latent tightly corresponds sources that have significant NMI values in that latent's s column. For example, either QLAE's 1 (first and second data samples) 27 (third data sample) causes changes in the camera height. The sensitivity of the NMI estimation can be observed in the low but non-negligible NMI values between the robot and robot y sources and the QLAE's Zg. Z3, Z,, Z5. Z7, Zy. Z11, Z13} manifesting as rare or small changes pose when intervening those latents. This sensitivity, however, means that it is important to remove inactive latents from InfoMEC estimation: the 3-TCVAE's 2 and Z11 are estimated have substantial NMI with background color and camera height, but do not actually cause any changes in the generations.
its range across annotated with source
22

---

QLAE
B-VAE
lighing innecitty
lightig innesiity
Fgghing
Fighing
Fighing
OUT
C:00
JUO
ECB
EC4
DE0
0.0
Fighing
600
OEI
030
C:31
RXC
AMO
GIT
COC
RCC
OUT
lighting
CRISETE
MXZ
DET
CRISETE
K:T:
XXY
CONTEC
40 A11
A11
Figure 14: Decoded Falcor3D latent interventions and NMI heatmaps for QLAE and 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. NMI heatmaps are annotated with source names, and inactive latents (as determined by range) have red font. In active latents, the visual variation the generations caused by each latent tightly corresponds sources that have significant NMI values in that latent's 's column. This holds for both obvious visual changes in sources like lighting and rather subtle visual changes in sources like camera x. Pruning inactive latents from NMI before computing InfoM and InfoC is particularly important for models that optimize for latent shrinkage, such as the V-V (InfoM, InfoC) is (0.71, 0.70) with pruning vs. (0.47, 0.51) without. Pruning is justified by the decoded interventions showing no dependence on the inactive latents.
23

---

QLAE
BioAE
stope
ibbox
Juc 1:1
Kon sox
our mbbox
bdl ous
mextt
mentt
obrect
son can
object scale
121
lighting intersity
lighting intersity
121
Sighing desation
lighiig disssiio
object coloe
son
obiect exioc
wall coloe
wall coloc
1121
Figure 15: Decoded Isaac3D latent interventions and NMI heatmaps for QLAE and 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 intervened on with a linear interpolation the dataset. The resulting representation is then decoded. NMI heatmaps are annotated with source names, and inactive latents (as determined by range) have red font. active latents, visual variation in the generations caused by each latent tightly corresponds to sources that have significant NMI values in that latent's 's column.
its range across
24

---

E
Quantitative Results
This section contains unabridged results for the experiments in the main text. Intervals denote 95% confidence intervals of the mean estimated assuming a t-distribution. Bolded intervals overlap with the interval with highest endpoint in the column. AE and InfoGAN variants are presented and bolded separately as all AEs are filtered for near-pertiee near-perfect data reconstruction, whereas InfoGANs are generally more lossy.
Table 12: Full results for the S dataset.
model
InfoM
InfoE
InfoC
D
If
C+
PSNR (dB) f
0.41 0.03
0.98 + 0.01
0.28 + 0.01
0.11
0.02
0.82 : 0.01
0.08
3.1
8-VAB
0.59 1: 0.02
0.99 1 0.02
0.49
: 0.03
0.61 0.00
0.99
1 0.01
0.47
1:47::00 0.01
3..1
B-TCAAE
0.98 1: 0.02
0.02
0.02
0.99
0.38
0.00
BioAE
0.56 4:6::000 0.02
0.01
0.44
0.02
0.33
0.00
0.25
31.6
VQ-VAE
0.97
0.07
0.81 0.10
0.31
36.0
AL (ours)
0.95 0.902 0.91
0.55
0.01
0.99
0.61
0.2
QLAE w/ global codebook
0.96 + 0.01
0.99
0.01
0.99
0.60
0.2
QLAE weight decay
0.99
0.02
0.03
0.3
VAE w/ weight decay
0.80 1 0.01
0.99 0.01
0.46
0.02
0.74
0.99
1 0.01
0.57 0.01
37.2
0.1
InfoWGAN-GP
0.00 0.03
0.41
0.02
0.23
0.00
0.80
0.01
0.18
12.0
QLInfoWGAN-GH (ours)
0.73 4: 0.03
0.76
0.38
0.00
0.85
0.01
0.29 2
20.4
QLInIoWGAN-GP w.d.
0.66
0.57
0.16
0.02
1.7
Table 13: Full results for the MPI3D dataset.
model
InfoM
InfoF
InfoC
D1
If
C1
PSNR (dB) 1
0.37 0.04
0.72
0.03
0.36 + 0.03
0.83 0.05
0.14
0.2
8-VAB
0.45 0.03
0.71 + 0.03
0.51 = 0.03
0.31 0.00
0.83 0.02
0.27 0.27
0.2
8-TCVAE
0.51
0.01
BioAE
0.45 0.00
0.66 0.00
0.36 0.03
0.24
0.79
0.19
37.9 1.22
VQ-VAE
0.57 0.06
0.22
0.00
0.03
0.04
35.3
0.2
QLAE (ours)
0.61
0.63
0.51
0.04
0.04
0.36
37.2
0.6
QLA w/ global codebook
0.54 1: 0.04
0.02
0.03
0.36
0.00
0.0
0.86
0.34
37.2 0.6
QLAE w/o weight decay
0.51
0.01
0.08
0.00
0.4 0.09
37.7
VQ-VAE w/ weight decay
0.50 + 0.04
0.01
0.01
0.22
0.60
0.20
34.7
0.2
InfoWGAN-GP
0.43 0.04
0.40 0.06
0.20
0.09
0.63
0.01 0.09
25.2 144
QI.nfoWGAN.co GP (surs)
0.62 0.03
0.51
+ 0.04
0.37
0.04
0.02
0.71 1 0.02
0.25
0.02
28.3
1.5
QLInfoWGAN-GP w/o w.d.
0.04
0.55
0.04
0.40
0.05
0.01
0.74
0.23
0.03
26.3
1.5
Table 14: Full results for the Falcor3D dataset.
model
InfoM
InfoE 1
InfoC 1
D1
11
C1
PSNR (dB) 1
AF
0.39 0.03
0.74 1: 0.03
0.20 1: 0.03
0.08
0.02
0.76
1: 0.02
0.07 00771000
0.01
29.9
3-VAE
0.71 0.05
0.73
0 0.04
0.70
0.03
0.32 0.02
0.84 + 0.02
0.28
29.0 0.2
8-TCVAE
0.66
0.71
0.71
0.00
0.33
0.1
BioAE
0.54 0.05
0.73 + 0.04
0.31
1:11:0000
0.81 0.00
0.02
2..
VQ-
0.05
0.42
0.79 0.00
0.29
29.0
QLAE (ours)
0.71 0.03
0.02
0.02
0.50
0.03
0.38
0.4
QLAE w/ global codebook
0.59 0.01
0.71
0.01
0.36
0.90
0.26
2.
AL w/o weight deray
0.65 - 0.02
0.76
0.43
0.06
0.97
1.02
0.36
0.44
VQ-VAE w/ weight doeay
0.74 0.02
0.86
0.40
0.02
0000
0.32
0.1
InfoWGAN-GP
1.44 0.60 0.03
0.30
0.04
0.11 0.01
0.08 0.01
1.69
QLInfoWGAN-GP (ours)
0.54
0.53
0:200 0.20
0.0
0.24
17.6
Q..nfoWGAN-GE GP w/a w.d.
0.56 0.08
0.56
0.04
0.52
0.02
0.04
0.72 0.02
0.07
1.1
25

---

Table 15: Full results for the Isaac3D dataset.
model
InfoM
InfoF. 1
InfoC
D1
11
C1
PSNR (dB) 1
AE
0.42 0.04
0.80 + 0.02
0.21 = 0.05
0.00
0.85
0.02
0.10 0.02
0.1
8-VAB
0.60
0.80 - 0.02
0.51
0.08
0.23
0.88
0.88
0.19
3.18
8-TCVAE
0.34 0.02
0.70 0.02
0.46
0.03
0.09
0.04
0.06
0.1
BioAE
0.63 1: 0.03
0.65
0.33 1 0.04
0.38
0.91
0.00
38.0
38.0
VQ-VAE
0.57 0.04
0.87
0.45
0.33
0.89
0.01
0.1
QLAE (ours)
0.78 + 0.03
0.19 0.07
00000000000000000000000000000000
0.09
0.99
0.02 0
0.3
QLAE global codebook
0.02
0.08
0.00 0.48
QLAE w/o weight decay
0.66 0.02
0.89
0.51
0.08
0.00
0.99
40.
VQ-VAE w/ weight decay
0.73 0.03
0.81
0.03
0.04
0.31
0.85
0.39
35.3 0.1
InfoWGAN-GP
0.53 + 0.03
0.51 0.02
0.24
0.71 0.13
0.11 21.9
QLInfOWGAN-Gi (ours)
000
0.25
2.7
QLInfoWGAN-GP w/o w.d.
0.59 0.01
0.53
00TE0000000000000000000001000000 0.23 2.12
26


Reviews:
[{"review_id": "Lcb2CihDVD", "reviewer": "Reviewer_sKhi", "summary": "This paper tackles the topic of learning disentangled representations from observations in a fully unsupervised manner. In particular, the authors propose a novel inductive bias for disentangled representation learning as well as an associated evaluation metric; The inductive bias relies on latent quantization, similar to the discretization of the latent space proposed in the VQ-VAE paper; Here authors propose an alternative motivated by the assumption that most scenes can we seen as a combination of distinct independent factors: each representation is an aggregation of latent factors each of which takes values from a unique codebook; Regarding the proposed metric for the evaluation of the ability of the aforementioned method to accurately disentangle sources of information, authors propose InfoMEC which encompasses three sub-scores for modularity, compactness, and explicitness measurement; This metric requires the availability of ground-truth sources of variation; \n\nThe authors propose the evaluation of their method using the InfoMEC metric on 4 synthetic benchmark datasets and show that in comparison to the $\\beta$-VAE, TC-VAE, and InfoGAN, the proposed QLAE outperforms benchmarks in terms of modularity and remains competitive in terms of the compactness metric; Additional ablations highlight that the success of the approach relies on the combination of the combinatorial discretization of the latent space and the heavy regularisation (high weight decay) applied during training; The simple adaptation of the VQ-VAE discretisation to the proposed latent quantization without heavy regularisation, leads to a modularity that is slightly above benchmarks;  ", "questions": "Questions: \n- Have authors tried to incorporate the proposed inductive bias into a Variational form of an AE? \n- Could authors elaborate on the datasets considered and in particular share statistical details regarding the ground-truth sources considered; \n- Could authors elaborate on why the bottleneck designed in the QLAE is stronger than the one in the VQ-VAE for a similar number of latent codes?\n- Have authors evaluated their method on more current benchmark evaluation metrics for disentanglement ?\n\nMinor comments: \n- Figures 3 & 4 are not referred to in the text \n- Typos in line 46, 50, caption Figure 2, 162, 265, 319\n- The granularity of the codebooks considered in the experiments is shared quite late in the text (discussion), this information would be relevant to share at an earlier stage; ", "limitations": "I agree with the limitations highlighted by the authors: \n\n- It remains unclear how this approach would scale to more complex datasets where underlying sources of information are continuous; This discretization of the latent representations is well motivated from an interpretability point of view. As highlighted by the authors the discrete nature of the representations also eases the evaluation; however, it remains unclear to me how one can motivate the use of the discrete representations in a generative setting; In real-world scenarios where data is more complex and arises from the mixing of a large set of continuous variables, how can the discrete bottleneck proposed by authors allow for the reconstruction of complex images; This limitation questions the impact of the contribution; \n\n- The contribution is centered around the combination of latent quantization and a heavy regularisation; Whist the intuition behind the first element seems clear to me, the paper does not provide any insights on why a heavy regularisation is a necessary condition to observe significant benefits; Whilst a contribution does not require an explanation behind each experimental showcased improvement, the absence of plausible explanation limits the impact of the paper by motivating the proposed approach by the experimental performance measurements solely; ", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- This paper is well written; Its content is easy to follow, the setup considered by the authors is clear and the experimental work whilst only encompassing a limited number of baselines, seems to have been rigorously conducted; The results shared by the authors allow the reader to easily grasp the main contributions and the ablations performed answer key questions the reader might have;\n\n- Whilst the methods seems to entail some limitations detailed below, the motivation for the proposed method (a combinatorial discretized latent space) considering the VQ-VAE results are sound; \n\n- The results shared by the authors highlight the positive contribution; The aggregated InfoMEC score over 4 datasets shows a positive improvement in terms of explicitness and modularity, two metrics that are sound to consider when evaluating disentanglement; By design, the approach is also more efficient than previous latent space discretization strategies by allowing a dimension-wise comparison for finding the nearest code to a continuous representation;", "weaknesses": "- I believe the authors could put more effort into precisely defining key concepts at earlier stages of the manuscript; The concept of compositional latent spaces, and parsimonious learning; \n- From my understanding, the field of disentangled representation learning encompasses a plethora of inductive biases and evaluation metrics (e.g., Eastwood et al. 2018: Disentanglement, Completeness, Information); Whilst authors define clearly the link with non-linear ICA, they remain more evasive regarding the link of the InfoMEC metric and previous metrics as well as the references to previous methods for disentangled representations; ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough evaluation of our submission and thought-provoking questions.\n\n> Have authors tried to incorporate the proposed inductive bias into a Variational form of an AE?\n\nOur proposed quantized latent autoencoder (QLAE) is variational to the same extent that VQ-VAE is, specifically with a uniform prior over discrete latent code vectors and a degenerate Dirac delta variational posterior. In our early experimentation, we did indeed try a variant that specifies a multinomial codebook posterior. We found that this didn’t confer any benefit and introduced further algorithmic complications, e.g. specifying a decaying schedule for the Gumbel-Softmax posterior temperature.\n\n> Could authors elaborate on the datasets considered and in particular share statistical details regarding the ground-truth sources considered;\n\nMPI3D consists of real images taken of a robotics apparatus, while the other 3 consist of simulation renders. We chose these datasets because they are the only ones we know of that feature i) high-dimensional, quasi-realistic observations and ii) more than 2 ground truth sources with complete labels for evaluation. \n\nEach dataset defines a set of 6 to 9 ground-truth sources, each of which can take on 4 to 40 discrete values. Both categorical and numerical sources are present. Each dataset is constructed from an exhaustive enumeration over all possible source vector realizations. This guarantees that each dataset’s empirical source distribution follows the assumption of mutually independent sources in the nonlinear ICA problem statement. \n\nWe have updated our manuscript to include source descriptions, source cardinality, and sample observations for each dataset in the Appendix.\n\n> Could authors elaborate on why the bottleneck designed in the QLAE is stronger than the one in the VQ-VAE for a similar number of latent codes?\n\nThis is a great question! Given the same number of latent codes and codebooks, there is no difference in terms of Shannon channel capacity, which is agnostic to the explicit form of each discrete code element. We conjecture that it is the maximally simplified choice of this form as a scalar for QLAE that better enables the encoder and decoder to attribute a consistent meaning to each value of each code element. However, we were not able to convincingly formalize this intuition, hence the empirical focus of the paper.\n\nWe have updated the figure contrasting vector quantization and latent quantization in our work to emphasize these points.\n\n> Whilst authors define clearly the link with non-linear ICA, they remain more evasive regarding the link of the InfoMEC metric and previous metrics as well as the references to previous methods for disentangled representations\n\nThank you for pointing this out. We address this in the \"Why use InfoMEC over previous disentanglement metrics?\" section of our main rebuttal.\n\n> Have authors evaluated their method on more current benchmark evaluation metrics for disentanglement ?\n\nYes, we have evaluated with random forest DCI. Please see the \"Quantitative results in terms of other disentanglement metrics\" section of our main rebuttal, as well as the rebuttal pdf for tables.\n\n> The authors propose the evaluation of their method using the InfoMEC metric on 4 synthetic benchmark datasets\n\nOne of the datasets (MPI3D) was collected on a real world robotics apparatus. While undoubtedly visually simplistic, it does contain some aspects of reality, e.g. noise.\n\n> Minor comments\n\nThank you for the pointers. We have added a reference to Figures 3 and 4 when comparing QLAE to the prior methods. We could not identify the typos you mention, but ran a spellcheck over the manuscript. We have added a statement regarding the fixed codebook size we use, as well as a pointer to a hyperparameters section in the Appendix, to the latent quantization section.\n\n> Whist the intuition behind the first element seems clear to me, the paper does not provide any insights on why a heavy regularisation is a necessary condition to observe significant benefits\n\nWeight decay has long been known to be a method of encouraging a very specific form of parsimony by “suppress[ing] any irrelevant components of the weight vector” [A]. Beyond this, we can only provide a handwavy intuition of how most solutions we would like to avoid (e.g. because they involve memorization) have a higher complexity than the solutions we want to find, motivating the use of weight decay to bias the search away from the former.\n\nMore recently, the use of weight decay has also taken a prominent role in the phenomenon of “grokking” [B], in which neural networks trained for extended periods, i.e. beyond overfitting, achieve strong extrapolation performance on small algorithmic tasks.\n\n> It remains unclear how this approach would scale to more complex datasets where underlying sources of information are continuous\n\nAs a first step towards answering this, we train QLAE on the CelebA dataset training split with 128 scalar latents and 10, 40, or 100 values per codebook. We compare QLAE to a vanilla AE and a $\\beta$-VAE, both with 128-dimensional continuous latent space. Overall, QLAE incurs a minor but significant penalty in terms of reconstruction. This penalty ameliorates only marginally with larger codebook size. The persistent discrepancy between AE and QLAE highlights the reduced capacity of combinatorial discretized latents compared to a continuous latent vector space.\n\n| model | val PSNR (dB) |\n|-|-|\n| AE | 21.2 |\n| $\\beta$-VAE | 20.5 |\n| QLAE, $n_v=10$ | 19.2 |\n| QLAE, $n_v=40$ | 19.5 |\n| QLAE, $n_v=100$ | 19.8 |\n\n#### References\n\n[A] Anders Krogh and John Hertz. \"A simple weight decay can improve generalization.\" NeurIPS, 1991.\n\n[B] Alethea Power, et al. \"Grokking: Generalization beyond overfitting on small algorithmic datasets.\" arXiv preprint arXiv:2201.02177, 2022.\n\n"}, {"review_id": "1zrBNi1MSi", "reviewer": "Reviewer_iB4p", "summary": "The authors propose to achieve disentanglement via latent quantization which constructs an inductive bias that matches the compositional nature of data. The authors also proposed three metrics to evaluate the disentanglement and achieved better performance according to the metrics they proposed. Ablation studies are conducted to understand different design choices.", "questions": "Please see questions in 'Weakness' section.", "limitations": "I did not realize potential negative societal impact.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "— A novel disentanglement method is proposed.\n\n— A few metrics on evaluating disentanglement are proposed to fix well-established shortcomings in previous metrics.\n\n— According to the experimental study on a few established datasets, the proposed methods dramatically improves the modularity and explicitness.", "weaknesses": "— The authors propose a few new metrics, and the metrics are sound. The authors claim that the metrics addressed a few issues of previous metrics, but experiments are missing to support the effectiveness of the metrics. \n\nIf the disentangled representations show very high InfoMEC, does it guarantee low reconstruction error, what does total correlation and MIG look like? If InfoMEC does not always align with total correlation, MIG and other existing disentanglement metrics, then diving deep into some concrete examples to explain why InfoMEC makes more sense.\n\n\n— The authors tried to do some qualitative measurements in Figure 4, but the authors only looked at QLAE and did not do a qualitative comparison with other methods.\n\n\n— Though the overall results in Table 1 looks amazing, the authors may elaborate more on a few things include:\n\n1. In terms of explicitness, it seems vanilla AE is doing pretty well or the best in most datasets. Does this mean the disentangled representation makes the relationship between sources and latents even less simpler?\n\n2. Also, according to the table, the proposed methods seem to significantly boost the representation learning with AE. But when the approach combined with InfoWGAN-GP, the boost is much smaller.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive evaluation of our work! We have taken your constructive criticism very seriously, and we believe we have addressed the most pressing points.\n\n> The authors propose a few new metrics, and the metrics are sound. The authors claim that the metrics addressed a few issues of previous metrics, but experiments are missing to support the effectiveness of the metrics.\n\nThank you for pointing this out. We appreciate that you find InfoMEC sound. We have included in the rebuttal pdf experiments demonstrating i) the sensitivity of nonlinear (random forest) DCI to tree depth; ii) the sensitivity of binning-based mutual information estimation to binning strategy; and iii) the robustness of the KSG mutual information estimator to the n_neighbors hyperparameter. Please see the \"Why use InfoMEC over previous disentanglement metrics?\" section of the main rebuttal to see how these experiments support our claims of improving over previous metrics. \n\n> diving deep into some concrete examples to explain why InfoMEC\n\nAside from the experiments above, we also include qualitative studies in the rebuttal pdf to show that the NMI matrix used in InfoMEC corresponds very tightly to the observed visual variation in the data when individual latents are intervened on (and then decoded). Please see the \"Additional qualitative results\" section of the main rebuttal and/or the Figure caption in the rebuttal pdf for more explanation.\n\n> The authors tried to do some qualitative measurements in Figure 4, but the authors only looked at QLAE and did not do a qualitative comparison with other methods.\n\nThank you for pointing this out. We have added to the Appendix qualitative comparisons between QLAE and the next-best prior method for every dataset. Due to space constraints, we could only include examples of this for one dataset (Isaac3D, which has the most sources among the datasets we considered) in the rebuttal pdf. We find that decoded latent interventions align very well with InfoMEC's NMI matrix. For QLAE, it is also worth pointing out that these generated samples are not reconstructions and involve decoding from continuous values not necessarily in the latent codebooks. Overall, the visual quality is reasonably high.\n\n> In terms of explicitness, it seems vanilla AE is doing pretty well or the best in most datasets. Does this mean the disentangled representation makes the relationship between sources and latents even less simpler?\n\nThis is a very interesting question. Adding latent quantization to vanilla AE does make the InfoE drop significantly for MPI3D, but for the other three, it maintains or increases the InfoE. It's not unreasonable that vanilla AE's explicitness is relatively high: we tune weight decay for vanilla AE to ensure fair comparison with QLAE. But yes, in the specific case of MPI3D, QLAE's representation is slightly less linearly predictive of the sources.\n\n> If the disentangled representations show very high InfoMEC, does it guarantee low reconstruction error?\n\nNo, it does not. This is why we filter the autoencoders for near-perfect reconstruction and InfoGANs for good reconstruction as part of our evaluation. We included reconstruction results in the supplementary material's Appendix C (apologies for the unfinished appendices in the submission pdf).\n\n> what does total correlation and MIG look like?\n\nThey would look very much like InfoC. All three measure compactness, i.e. the extent to which latents encode information about disjoint sets of sources. InfoC is essentially MIG using a better mutual information estimator and a different aggregation (ratio of max to sum) in order to not leave out most entries of the NMI matrix in the computation, as the gap aggregation does. \n\n> Also, according to the table, the proposed methods seem to significantly boost the representation learning with AE. But when the approach combined with InfoWGAN-GP, the boost is much smaller.\n\nThis is true, but note that the InfoWGAN-GP has a higher InfoM than vanilla AE, resulting in slightly less room for improvement. We also note that autoencoder variants are much more widely used in disentangled representation learning, so our experiments with InfoGAN were more intended to rule out the possibility that latent quantization was exploiting something about autoencoders.  But overall we agree that latent quantization's benefit for InfoGAN is less dramatic than its improvement for vanilla AE. \n"}, {"review_id": "KgJ2OKoMW8", "reviewer": "Reviewer_hQMG", "summary": "The paper discusses the problem of disentangled representation learning (finding latent variables only from observed data that ideally correspond to the true underlying sources of variation in the data). The authors suggest both a method for disentangled representation learning and a metric for evaluating it. The suggested method, termed latent quantization (QLAE), is based on finding an element-wise quantized representation (similar to VQ-VAE, but with per-dimension quantization). The quantized representation is learned using SGD in a similar manner to VQ-VAE. The proposed metric (InfoMEC) is divided to three values - Modality, Explicitness and Compactness. These values are computed using mutual information between the true latents and the discovered ones, with different normalizations. The authors then demonstrate that QLAE achieves the best InfoMEC scores using several datasets commonly used for disentangled representation learning.", "questions": "Does element-wise quantization work well when the true latent factors are continuous?\n\nIn the experiments, images were downscaled to 64x64. Can this affect the evaluation?\n\nThere is no reference to the figures in the text itself (e.g. figure 1 shows...). This makes it a bit harder to understand the figures in the right context.\n\n", "limitations": "The authors discuss two current limitations of the paper: \n- The lack of analytical proof: \"While currently we can only support this with empirical evidence, our key motivation is that generative models for realistic data are compositional, so our latent space should be too.\"\n- A possible \"overfitting\" of the proposed latent quantization methods to the benchmark datasets \"in which, just like our model, the sources are discrete and the generative process is noiseless\"\n\nAnother limitation, which I'm not sure is discussed, is the possible bias between the proposed representation learning method and the proposed evaluation method - QLAE is only evaluated (and compared to prior methods) using InfoMEC.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper is clearly written - all mathematical notations are well defined and both the representation learning and evaluation methods are clearly motivated and defined.\n\nThe contribution of a InfoMEC, measuring modularity, explicitness and compactness could be significant to the disentangled representation learning community, which perhaps lacks well-motivated and accepted evaluation methods.\n", "weaknesses": "The paper is proposing both a disentangled representation learning method and a method to evaluate it. Perhaps due to lack of space, I think that this results in insufficient discussion of these two separate contributions. Specifically:\n\n- The prior work section does not discuss existing metrics for disentanglement nor does it compare InfoMEC to them or demonstrate the advantages of InfoMEC. For example, compared to Mean Correlation Coefficient (MCC) [1]\n- The proposed method is only evaluated using InfoMEC. This is a significant drawback in my opinion.\n- An example for prior work in disentangled representation learning that is not discussed is IRMAE [2], where the authors encourage an autoencoder to learn a disentangled representation using rank minimalization.\n- The relationship between compositionality and element-wise quantization is demonstrated empirically, but the intuition why element-wise quantization encourages (or discovers) compositionality was not clear enough to me and no mathematical proof is provided.\n\n[1] Khemakhem et al, Variational Autoencoders and Nonlinear ICA: A Unifying Framework, AISTATS 2020\n\n[2] Jing et al, Implicit Rank-Minimizing Autoencoder, NeurIPS 2020", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and appreciation of the strengths of our work. We have striven to address your main concerns with our submission.\n\n> The proposed method is only evaluated using InfoMEC. This is a significant drawback in my opinion. [...] Another limitation, which I'm not sure is discussed, is the possible bias between the proposed representation learning method and the proposed evaluation method - QLAE is only evaluated (and compared to prior methods) using InfoMEC.\n\nWe have repeated our entire evaluation protocol using random forest DCI instead of InfoMEC. Tables of both are included in the rebuttal pdf. Overall, the DCI results are in high agreement with the InfoMEC results, especially for modularity. Please see the \"Quantitative results in terms of other disentanglement metrics\" section of the main rebuttal for details on our implementation of DCI, as well as Pearson and Spearman correlation tests between corresponding members of InfoMEC and DCI.\n\n> The prior work section does not discuss existing metrics for disentanglement nor does it compare InfoMEC to them or demonstrate the advantages of InfoMEC. For example, compared to Mean Correlation Coefficient (MCC)\n\nThank you for pointing this out. We have added explicit discussion of existing metrics (including MCC) and how InfoMEC improves upon them to the InfoMEC section of the manuscript. Please see the \"Why use InfoMEC over previous disentanglement metrics?\" section in the main rebuttal for a self-contained exposition. We also provide a demonstration of the advantages of our use of the KSG mutual information estimator (highly increased robustness to estimation hyperparameters) to the Appendix. This is included in the rebuttal pdf.\n\n> An example for prior work in disentangled representation learning that is not discussed is IRMAE [2], where the authors encourage an autoencoder to learn a disentangled representation using rank minimalization.\n\nThank you for the reference. We have added a discussion of ideas that are conceptually compatible with latent quantization, including IRMAE, to our manuscript.\n\n> The relationship between compositionality and element-wise quantization is demonstrated empirically, but the intuition why element-wise quantization encourages (or discovers) compositionality was not clear enough to me and no mathematical proof is provided.\n\nWe apologize for the confusion. As Reviewer XVcp rightly points out, we should not be characterizing our latent space or models as \"compositional\", as we simply have no evidence for such claims. We have revised our manuscript to avoid using \"compositional\" for anything other than true generative processes.\n\n> Does element-wise quantization work well when the true latent factors are continuous?\n\nAs a first step towards answering this, we train QLAE on the CelebA dataset training split with 128 scalar latents and 10, 40, or 100 values per codebook. We compare QLAE to a vanilla AE and a $\\beta$-VAE, both with 128-dimensional continuous latent space. Overall, QLAE incurs a minor but significant penalty in terms of reconstruction. This penalty ameliorates only marginally with larger codebook size. The persistent discrepancy between AE and QLAE highlights the reduced capacity of combinatorial discretized latents compared to a continuous latent vector space.\n\n| model | val PSNR (dB) |\n|-|-|\n| AE | 21.2 |\n| $\\beta$-VAE | 20.5 |\n| QLAE, $n_v=10$ | 19.2 |\n| QLAE, $n_v=40$ | 19.5 |\n| QLAE, $n_v=100$ | 19.8 |\n\n> There is no reference to the figures in the text itself (e.g. figure 1 shows...). This makes it a bit harder to understand the figures in the right context.\n\nThank you for pointing this out. We have ensured that all figures and tables have a reference to them from the main text.\n\n> In the experiments, images were downscaled to 64x64. Can this affect the evaluation?\n\nThis certainly makes the modeling significantly easier for the Falcor3D and Isaac3D datasets (Shapes3D and MPI3D are natively 64 x 64). However, it should not make disentangling easier. We valued the unification of experimental procedure enabled by downsampling."}, {"review_id": "tCmEzf6bja", "reviewer": "Reviewer_XVcp", "summary": "The paper presents an alternative approach to the standard disentanglement representations learning framework whereby the continuous latent space is substituted by a quantized one. Based on this new quantized latent space, the authors propose novel ways to measuring several properties related to disentanglement. They then make two main claims: 1) that this quantized latent space better aligns with how generative factors are usually encountered in data (both real and synthetic) and 2) that their novel disentangled measures more robustly measure the real disentanglement learned by the model.", "questions": "For the first point, the authors claimed to have tested 5 datasets and two variants of their model against some previous models. However, the scores for their framework are only provided for the 3DShapes dataset. While it is true that space is limited, I would have expected them to include the other plots in the appendix. Similarly, I would have expected to see more traversal reconstructions for the other datasets in the appendix. Speaking of which the authors left a comment in big red bright letters which I assume was for their own benefit. Always review before submitting!\n\nI assume including these plots in the appendix will be easy enough. What I do expect to see is a visualization of the latent space of the model (even if it is only for arbitrary pairs of variables). They clearly visualize how the latent space should look for their model and then… proceed to not show how it actually looks for trained models. This should have been one of their core plots. Yes, metrics are better for comparing a large number of models, but visual inspection can offer important insights and one of their core contributions is that a quantized latent space should be more transparent.\n\nThe second issue is that their InfoMEC framawork is very similar to the DCI framework of Eastwood & Williams (2019), whose concepts of disentanglement and completeness are equivalent to those compactness and modularity, respectively. A comparison between the two is thus warranted to prove the effectiveness of their approach. This not only requires overall scores but NMI visualizations and visualizations of the latent space in order to check that the authors claims (i.e. that their metrics are better) are indeed true.", "limitations": "There is also a lack of discussion on some of the limitations of this approach. One reason why a continuous space is often preferred is because it is not necessary to set any other parameters. While I agree with the authors hypothesis that a quantized space should be better, it also introduced the problem of setting the size of this quantized space. What if we do not know how many sources there are? The authors need to at least discuss how this is a limitation of their proposal.\n\nFinally, there is a recurring conceptual issue in the introduction, namely the conflation of disentanglement (which is what they authors explore) with compositionality (which they are not). While it is certainly possible to take inspiration from how generative factors in disentanglement datasets are defined (and this is indeed what the authors do), the fact that those representations are compositional does not entail that the quantized space that the authors propose will also be compositional.\n\nTo see why notice that the definition of compositionality is that a system uses its knowledge of existing entities and rules of composition of a particular domain to understand novel instances (see [1]). Thus compositionality is a mechanism to transfer performance from training to out-of-distribution settings. All the experiments in the current work are in-distribution and thus do not test for compositionality.\n\nIn the case of the ground truth generative factors, these are compositional because they are used by very specific image generators that understand the semantics of these values in the context of the images they must generate. Whether the decoder used by VAEs has this capacity (which would make the representations compositional) is something that the authors do not test (and is actually not the case! See below).\n\nI do not expect the authors to explore the compositionality issue here as this would be an entirely different article. But they should be able to partly rewrite the introduction accordingly, limiting their references to compositionality to the very minimal (we take inspiration from…) and refrain from referring to the quantized latent space as compositional. \n\nIf the authors are interested I would recommend Duan et al., for a hypothesis of the relations between disentanglement and compositionality and subsequent work in Montero et al., Schott et al.. These show that disentanglement does not necessarily lead to compositionality. Of particular interest to the authors is the experiments in Montero et al, 2021 which show that having a quantized and disentangled latent space (i.e. the ground truth one) does not lead to compositionality in reconstruction. Rozell et al do use quantized representations but their architecture is very different from that of a standard VAE (and the one the authors propose).\n\nIn any case, taking into account these issues I cannot currently recommend acceptance of this article. I will be willing to revise my score should these points be addressed.\n\n\nReferences:\n\n[1] Duan, S., Matthey, L., Saraiva, A., Watters, N., Burgess, C. P., Lerchner, A., & Higgins, I. (2019). Unsupervised model selection for variational disentangled representation learning. arXiv preprint arXiv:1905.12614.\n\n[2] Montero, Milton Llera, et al. \"The role of disentanglement in generalisation.\" International Conference on Learning Representations. 2020.\n\n[3] Schott, L., Von Kügelgen, J., Träuble, F., Gehler, P., Russell, C., Bethge, M., ... & Brendel, W. (2021). Visual representation learning does not generalize strongly within the same domain. arXiv preprint arXiv:2107.08221.\n\n[4] Montero, M. L., Bowers, J. S., Costa, R. P., Ludwig, C. J., & Malhotra, G. (2022). Lost in latent space: Disentangled models and the challenge of combinatorial generalisation. arXiv preprint arXiv:2204.02283.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "In exploring the idea of using a quantised latent space to learn a disentangled representation of the data, the paper has several key strengths:\n1. The authors use a good theoretical motivation based on results from the ICA literature.\n2. The observation that disentangled datasets use quantized generative factors (even if the values are intended to be over the reals) is one that, even if obvious, not many researchers have explored (though some have, see below).\n3. The disentanglement metrics are well motivated based on this quantized representation.\n4. The authors used two main variants and several datasets to test their approach which gives their results robustness.", "weaknesses": "On the other hand there are several places where the paper as is currently written is lacking in my opinion. In no particular order these are:\n\n1. While the authors present numerical results for five datasets, they only provide traversal reconstructions for one, visualization of one metric for one dataset and no visualization of the actual latent values.\n2. Despite the fact that the proposed disentanglement measurement framework bears a significant resemblance to the one proposed in Eastwood and Williams, they only reference this prior work in passing without giving any in depth comparison between the two.\n3. The authors make the point several times that they are using a compositional representation (e.g. with phrases such as “*A* *side advantage of the compositional latent encoding…”)* yet they seem to misunderstand what compositionality is and whether their experiments test for this.\n4. There is no discussion of any limitations of their approach even though some of them are obvious (like the fact that we need to define the size of the quantization).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your extensive review. We agree with the main criticisms you have about our manuscript at the time of submission, and we have taken action to address them. We appreciate your willingness to re-assess our work.\n\n> While the authors present numerical results for five datasets, they only provide traversal reconstructions for one, visualization of one metric for one dataset and no visualization of the actual latent values. [...] What I do expect to see is a visualization of the latent space of the model [...]  visual inspection can offer important insights and one of their core contributions is that a quantized latent space should be more transparent.\n\nWe agree with all of these points. We have added decodings of latent interventions for all datasets considered to the Appendix, and include examples of this for QLAE and BioAE on Isaac3D in the rebuttal pdf. Please see the \"Additional qualitative results\" section of the main rebuttal for more explanation.\n\n> A comparison between [InfoMEC and DCI] is thus warranted to prove the effectiveness of their approach. This not only requires overall scores ...\n\nWe have repeated our evaluation protocol using random forest DCI instead of InfoMEC, and have included tables of these results in the rebuttal pdf. QLAE consistently scores highest in modularity in terms of both InfoM and D. Please see the \"Quantitative results in terms of other disentanglement metrics\" section of the main rebuttal for details on our implementation of DCI, as well as Pearson and Spearman correlation tests between corresponding members of InfoMEC and DCI.\n\n> Despite the fact that the proposed disentanglement measurement framework bears a significant resemblance to the one proposed in Eastwood and Williams, they only reference this prior work in passing without giving any in depth comparison between the two.\n\nWe apologize for neglecting this in our submission. We have added explicit discussion of DCI (linear and nonlinear) to the InfoMEC section of our manuscript. Please see the \"Why use InfoMEC over previous disentanglement metrics?\" section of the main rebuttal, in particular paragraph 3, for a self-contained explanation of how InfoMEC avoids key drawbacks inherent to nonlinear (random forest) DCI. We also provide some experimental evidence supporting these claims in the rebuttal pdf. \n\n> However, the scores for their framework are only provided for the 3DShapes dataset [...] include the other plots in the appendix.\n\nWe have added NMI heatmaps for QLAE and the best prior method for all datasets to the Appendix. Two examples are included in the rebuttal pdf.\n\n> The authors make the point several times that they are using a compositional representation (e.g. with phrases such as “A side advantage of the compositional latent encoding…”) yet they seem to misunderstand what compositionality is and whether their experiments test for this.\n\nWe agree that we should not use \"compositional\" to describe our latent space design. We have removed all mention of \"compositional\" from the paper, aside from when describing true generative processes. We instead use words like \"organized\" and \"combinatorial\" to describe the latent space.\n\n> There is no discussion of any limitations of their approach even though some of them are obvious (like the fact that we need to define the size of the quantization).\n\nWe did include discussion of our limitations, including a mention of the quantization size, in the Discussion section in our submission, but we apologize for not making this information more readily apparent. We have restructured this section with an explicit Limitations paragraph heading. \n\n> Speaking of which the authors left a comment in big red bright letters which I assume was for their own benefit. Always review before submitting!\n\nWe apologize for this oversight. The appendices included in our supplementary material are much more complete.\n"}, {"review_id": "cDWq5bL03o", "reviewer": "Reviewer_Mf8i", "summary": "The paper proposes learning a disentangled latent representation by quantizing the latent space into learnable discrete values, with a separate scalar codebook for each dimension. The main assumption that the paper relies on is that \"many datasets of interest are generated from their sources in a compositional manner.\" Additionally, the paper introduces a new metric to better evaluate the disentanglement.", "questions": "1) In line 49, the author mentions that \"they enable the use of simpler, more robust, distribution estimation techniques than continuous latent spaces.\" Could you elaborate on that?\n\n2) The new interpretation given in equation 3 for infoGAN does not appear convincing.\n\n3) In line 121, it is mentioned that \"generative models for realistic data are compositional.\" What is meant by generative models being compositional? Earlier, it was mentioned that realistic data is generated in a compositional manner. I assume these two statements imply the same thing. I wonder if this assumption oversimplifies the nature of generative models and may not be generalizable.", "limitations": "The limitation of the paper is discussed and there is no potential negative societal impact of their work. ", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper proposed a new metric for disentanglement.\n", "weaknesses": "\n\n1) The manuscript could benefit from improvements as some parts and notations are rather confusing.\n\n2) The proposed method appears to be very similar to VQ-VAE.\n\n3) It is uncertain if the underlying assumption that real-world data is generated in a compositional manner holds.\n\n4) Concerning disentanglement, the paper only presents results on the metric they propose. It would be interesting to see a comparison with VQ-VAE using other disentanglement metrics.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We believe we can address the main concerns you had about our submission, in particular the similarity to VQ-VAE and only evaluating with InfoMEC.\n\n> The proposed method appears to be very similar to VQ-VAE.\n\nIndeed, latent quantization can be summarized as three modifications to vector quantization: i) scalar instead of vector codes, ii) atypically high weight decay on the encoder and decoder, and iii) separate per-dimension codebooks instead of a single global codebook. Each of these changes is ablated with our full experimental protocol on autoencoders and all are shown to be necessary for best performance (Table 5). Even if one views these contributions as not constituting much algorithmic novelty, we feel there is substantial impact in identifying and verifying a recipe for dramatic improvement on a difficult problem. In any case, we take care to make proper attributions to the VQ-VAE work in the technical section, algorithm pseudocode, and related works section. We have also refined our Figure 2 to more clearly delineate the differences between vector quantization and latent quantization.\n\n> It is uncertain if the underlying assumption that real-world data is generated in a compositional manner holds.\n\nWhile there may exist datasets that aren’t accurately described by this assumption, it certainly holds for the ones we (and the field of disentangled representation learning) consider. For example, take the “complex shapes” variant of MPI3D we use, which consists of images taken of a real world robotics apparatus. Changing any particular source dimension while holding the others fixed corresponds to a consistent and modular change in the data, e.g. via physically actuating the manipulator, swapping out the object, or moving the camera.\n\n> Concerning disentanglement, the paper only presents results on the metric they propose. It would be interesting to see a comparison with VQ-VAE using other disentanglement metrics.\n\nThank you for pointing this out. We have now added random forest DCI evaluation. Please see the \"Quantitative results in terms of other disentanglement metrics\" section of our main rebuttal, as well as the rebuttal pdf for tables.\n\nIn particular, the improvement of QLAE over VQ-VAE holds under DCI, too.\n\n> In line 49, the author mentions that \"they enable the use of simpler, more robust, distribution estimation techniques than continuous latent spaces.\" Could you elaborate on that?\n\nWhen a source and latent are both discrete, their joint distribution can be characterized by explicitly constructing a table of frequencies over an evaluation dataset. This is maximally expressive, so we are freed from needing to make any approximating modeling choices, such as binning strategy, number of nearest neighbors, kernel, or variational family. We briefly mention this advantage in the context of mutual information estimation in Section 4.1 (specifically line 140). \n\n> The new interpretation given in equation 3 for infoGAN does not appear convincing.\n\nThe interpretation of the InfoGAN encoder loss as latent reconstruction serves to justify practical implementations (e.g. mean-squared error) as natural design choices rather than loose variational bounds on mutual information. The interpretation also provides a lens with which to see autoencoders and InfoGANs as duals, motivating our demonstration of the applicability of latent quantization to both in our experiments. Since these points are entirely tangential to the submission’s main contributions, we opted for a very condensed statement; our apologies if this detracted from the reading experience. Note also that the interpretation aligns closely with the method pseudocode (Algorithm 3, Appendix A).\n\n> “generative models for realistic data are compositional”\n\nThank you for pointing this out. There is a typo of “generative models”, which we have amended to “generative processes”, i.e. the real world (or high-fidelity simulation) physics that generate the data. We hope this clears up the confusion.\n\n> The manuscript could benefit from improvements as some parts and notations are rather confusing.\n\nIf there is anything specific beyond what you already listed as questions, please do let us know, as your constructive feedback is much appreciated.\n"}]

Author rebuttal:
We thank the reviewers for their detailed and insightful feedback. Our manuscript has greatly benefited as a result of addressing the reviewers’ most critical comments. A brief summary of the major changes:
- **Quantitative results based on DCI evaluation [all reviewers].** We have repeated our evaluation procedure substituting DCI for InfoMEC, and have added the results to the Appendix. 
- **Explanation of InfoMEC’s advantages over previous metrics [sKhi, iB4p, hQMG, XVcp].** We have rewritten the InfoMEC section to include explicit discussion with linear DCI, nonlinear DCI, MIG, MCC, and SAP, highlighting the advantages of InfoMEC over each. Experiments supporting our main points have been added to the Appendix.
- **Additional qualitative results [iB4p, XVcp]**. We have added decoded latent traversals and NMI heatmaps for QLAE and the prior method with second-highest InfoM for every dataset to the Appendix. These are pointed to from the experiments section and relevant figure captions. 

We now respond to the main concerns raised by the reviewers.

### Quantitative results in terms of other disentanglement metrics [all reviewers].
We run random forest DCI with a 0.9/0.1 train/val split of 10K (source, latent) samples, tuning the max tree depth hyperparameter w.r.t. val accuracy. **We include random forest DCI results in the rebuttal pdf.** We also include InfoMEC results for comparison.
QLAE consistently improves the modularity of representations over prior works in terms of InfoMEC and DCI. The scaling of DCI actually makes the improvement seem bigger. Each pair of metrics measuring the same (or similar) property are correlated with vanishing p-values. The correlation is especially high for the modularity metrics.

| metrics pair | Spearman $\rho$ | Pearson $r$ |
| - | - | - |
|  (InfoM, D) | 0.82 | 0.80 |
| (InfoE, I) | 0.77 | 0.76 |
| (InfoC, C) | 0.59 | 0.57 |

### Why use InfoMEC over previous disentanglement metrics? [sKhi, iB4p, hQMG, XVcp]
Our desire to formulate InfoMEC stemmed from our frustration in working with prior disentanglement metrics like DCI and MIG, combined with calls for improvements from survey efforts [9]. 

InfoMEC's use of mutual information to quantify the relatedness between an individual source and latent aligns closely with nonlinear ICA, which allows *any* invertible transformation between them. In contrast, measures such as correlation (used in MCC), relative LASSO weights (used in linear DCI), or linear predictive accuracy (used in SAP) fail to detect many allowed functional relations. 

Nonlinear DCI [16] suffers several practical drawbacks from being defined in terms of relative counts of decision tree splits: determining this requires considering all latents *jointly* while fitting (all latents) -> (one source) trees. This results in a cumbersome computational footprint that is exacerbated by highly sensitive hyperparameters such as maximum tree depth [9], the tuning of which has even been omitted in prior work [39]. 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 (all latents) -> one source mappings, but does so with function classes of severely limited capacity for which fitting procedures scale well.)  **See the rebuttal pdf for a plot showing DCI's sensitivity to the maximum tree depth hyperparameter.**

InfoM and InfoC are essentially improved versions of MIG-sup [38] and MIG [10], following the strategy of i) pairwise source-latent mutual information estimation followed by ii) aggregation. Our improvements fix obvious issues in these steps identified by [9]. Firstly, common implementations of MIG and variants use histogram binning to approximate the marginal distribution of a single continuous latent. (The original MIG work [10] used a bespoke estimator for models with probabilistic latents, which cannot be used for models such as the vanilla AE, BioAE, VQ-VAE, or QLAE.) We fix this by prescribing the use of the celebrated $k$-neighbors based KSG estimator and variants, which are very robust with respect to $k$ since the variables involved are one-dimensional. **See the rebuttal pdf for plots showing histogram binning sensitivity to binning strategy and KSG robustness to $k$.** Secondly, the choice of gap as the aggregation function results in an indifference to most of the values in the pairwise normalized mutual information matrix. We fix this by prescribing the use of the ratio aggregation introduced by [59]. Compared to “mutual information gap/ratio” and “mutual information gap supplement”, our proposed names make it much clearer which disentanglement property each measures. Finally, our formalization of InfoE using the predictive $\mathcal{V}$-information framework both unifies it with InfoM and InfoC as information-theoretic and exposes the predictive family as a design choice.

### Additional qualitative results [iB4p, XVcp]
We have added decoded latent traversals and NMI heatmaps for QLAE for every dataset in the Appendix. For comparison, we use the prior method that produced the second-highest modularity for each dataset. We remark that our QLAE traversal is a linspace from the minimal to the maximal value of that latent (rather than an enumeration of the learned discretized values), verifying QLAE’s interpolative generalization capabilities. We find that the decoded latent traversals and NMI are in high agreement for both QLAE and the prior method. In particular, one can observe low but nonzero NMI$[i,j]$ value corresponding to source $i$ being sensitive to changes in latent $j$ only for some values or some samples. **We have included an example of these qualitative results for the Isaac3D dataset in the rebuttal pdf.**

Thank you again for your time and energy in reviewing our work!





Meta-review: The reviewers raised several concerns about the paper, but the authors have addressed all of these concerns in their rebuttal. The authors have also added new quantitative and qualitative results that support their claims.

The new quantitative results show that InfoMEC is a better measure of disentanglement than previous metrics, such as DCI and MIG. The authors also show that QLAE consistently outperforms previous VAE methods on a variety of datasets.

The new qualitative results provide further evidence of the disentangling ability of QLAE. The authors show decoded latent traversals and NMI heatmaps for QLAE and the prior method with the second-highest modularity for each dataset. These results show that QLAE is able to learn disentangled representations that are more interpretable than those learned by previous methods.

Overall, the paper makes a significant contribution to the field of disentangled representation learning. The authors have proposed a new method for quantifying disentanglement that is more reliable than previous metrics. They have also introduced a new VAE model that is specifically designed to improve disentanglement. The new quantitative and qualitative results provide strong evidence of the effectiveness of InfoMEC and QLAE.

I recommend that this paper be accepted for publication.
Decision: Accept (poster)