Title: Disentanglement via Latent Quantization
Venue: NeurIPS 2023
Accepted: yes
Keywords: disentanglement

unsupervised learning

quantization
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

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.

simple

complex

v.s.

??

??

(a)

(b)

(c)

(d)

(e)

1 Introduction

Our increasing reliance on black-box methods for processing high-dimensional data underscores the importance of developing techniques for learning human-interpretable representations. To name but a few possible benefits, such representations could foster more informed human decisionmaking [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 a 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 a need for an inductive bias that enables consistently good performance in the purely unsupervised setting.

In this work, we answer the call for a better inductive bias for disentanglement. Our solution is motivated 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, our broad strategy to 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 a small number of scalar values to 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.

A 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 datareconstructing (vanilla autoencoder) and latent-reconstructing (InfoGAN) generative models. Together 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-truth source evaluations [9, 20, 53]. In particular, our quantized-latent 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. 1

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-generating model of nonlinear independent components analysis (ICA) [30], a problem very related to but more conceptually precise than disentanglement:

where s = ( s 1 , . . . , s n s ) comprise the n s 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 a dataset D of samples from this model. Specifically, the solution should include an approximate inverting function ˆ g -1 : X → Z such that, assuming latent variables (latents) z = ( z 1 , . . . , z n s ) are matched correctly to the sources, each source is perfectly determined by its corresponding latent. A typical technical phrasing is that ˆ g -1 ◦ g should be the composition of a permutation and a 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 a 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 it under sufficiently generalizable assumptions 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 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 following three properties: modularity -the extent to which sources are encoded into disjoint sets of latents; explicitness -how 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 a cohesive information-theoretic framework that we name InfoMEC.

The disentangled representation learning problem statement considered in this work is as follows. Given a dataset of paired source-data samples { ( s, x = g ( s )) } from the nonlinear ICA model (1), learn an encoder ˆ g -1 : X → Z and decoder ˆ g : Z → X solely using the data { x } such that (i) the InfoMEC as estimated from samples { ( s, z = ˆ g -1 ◦ g ( s )) } from the joint source-latent distribution is high, while (ii) maintaining an acceptable level of reconstruction error between x and ˆ g ◦ ˆ g -1 ( x ) .

2.2 Autoencoding and InfoGAN as Data and Latent Reconstruction

We will apply our proposed latent quantization scheme to two foundational approaches for disentangled representation learning: vanilla autoencoders (AEs) and information-maximizing generative adversarial networks (InfoGANs) [12]. Here, we provide a brief overview of the two and defer complete implementation details to Appendices A and C. Both approaches involve learning an encoder ˆ g -1 and decoder ˆ g . An autoencoder takes a datapoint x ∈ X as input and produces a reconstruction ˆ g ◦ ˆ g -1 ( x ) ∈ X that is optimized to match the input:

An InfoGAN instead takes a latent code z ∈ Z as input. The decoder (aka generator) maps z to the data space, and from this the encoder produces a reconstruction of the latent:

L

reconstruct latent

(ˆ

g

-

1

,

ˆ

g

) :=

E

z

∼

p

(

z

)

[

-

log

p

(

z

|

ˆ

g

-

1

◦

ˆ

g

(

z

))

]

.

(3)

Unlike the data reconstruction loss, this is clearly insufficient for learning as the dataset D isn't even used. InfoGAN can be thought of as grounding latent reconstruction by making the marginal distribution of generated datapoints, ˆ g ( 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 nonidentifiability of nonlinear ICA? Our key motivation is that

(a) vector quantization

(b) latent quantization

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, we propose latent quantization, a modification of vector quantization [55]. In the latter, the latent representation of a datapoint is partitioned into component vectors, each of which is quantized to the nearest of a discrete set 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 n z distinct scalar codebooks: Z = V 1 ×··· × V n z , where each codebook V j is a set of n v reals. The complete transformation comprises, first, the encoder network mapping the data to a continuous latent space ˆ g -1 : X → R n z (with a slight abuse of notation), followed by quantization onto the nearest code (Figure 2b). This nearest neighbor calculation can be done elementwise, which is highly efficient. Formally, latent quantization is:

We represent Z by its constituent discrete values, concretely as a learnable two-dimensional array V ∈ R n z × n v , the j -th row of which stores the elements of codebook V j in an arbitrary order.

A lesson from nonlinear ICA is that, given a flexible enough model, data can be mapped to and from latent spaces in 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 or something close enough to it. We operationalize this by using a 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.

To train a quantized-latent model, we use the straight-through gradient estimator [5] and co-opt the quantization and commitment losses proposed for vector quantization:

The straight-through gradient estimator facilitates the flow of gradients through the nondifferentiable quantization step. L quantize pulls the discrete values constituting z onesidedly towards the prequantized continuous output of the encoder. This is 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 a 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 a 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 1 contains pseudocode for latent quantization and computing the quantization and

commitment losses. Appendix A presents pseudocode for training a quantized-latent autoencoder (QLAE) in Algorithm 2 and a quantized-latent InfoWGAN-GP [12, 1, 23] in Algorithm 3.

4 InfoMEC: Information-Theoretic 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 to 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 to recover the sources up to a permutation and dimension-wise invertible transformation. The mutual information between an individual source and latent,

is a 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 non-trivial. Previous works bin a continuous variable and pretend it is discrete [47, 10], but this is sensitive to the binning strategy [10]. Instead, for evaluating continuous latents, we choose the celebrated k -nearest neighbor based KSG estimator [40, 19], in particular a variant [60] designed to handle a 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 k (Figure 6). We remark that latent quantization enables reliable evaluation using the discrete-discrete estimator.

To facilitate aggregation, we desire a normalization to the interval [0 , 1] . To this end, note that the identity I ( s i ; z j ) = H ( s i ) -H ( s i | z j ) and the nonnegativity of entropy imply I ( s i ; z j ) ≤ H ( s i ) for discrete sources. Following [11], we define a normalized mutual information as

We prefer this normalization scheme over others [66] since i) it is the proportion of a source's entropy reduced by conditioning on a latent and thus scales consistently to [0 , 1] for any model, and ii) it avoids the scale-dependent (and possibly negative) differential entropy of a continuous latent. We gather all evaluations of NMI( s i , z j ) into a 2-dimensional array NMI ∈ [0 , 1] n s × n z . We remove inactive latents (columns of NMI ), which are those with zero range (over the evaluation sample) for discrete latents. For continuous latents, zero is too strict, so we heuristically define the threshold to be 1 /20, applied after dividing the ranges by their maximum. See Figure 3 for examples of NMI ⊤ .

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 NMI has only one nonzero element. This has been measured as the gap [11] between the two largest entries in a

AE

QLAE

-TCVAE

20

0.08

0.03

0.06

0.05

0.02

10'0

0.04

0.13

0.01

0.00

0.00

80'0

2

0.03

0.02

0.01

0.04

0.02

0.21

0.04

0.00

0.01

0.28

0.03

10'0

0.00

0.03

0.18

0.01

0.01

90'0

0.04

0.01

0.01

L00

0.22

0.02

0.15

0.00

0.01

0.07

0.05

0.01

0.03

0.03

0.05

t0'0

0.01

0.10

0.05

0.05

0.01

0.02

0.02

0.13

0.01

0.12

0.00

0.01

0.25

0.02

olz

0.02

0.04

0.07

0.00

0.03

0.09

11Z

0.14

0.07

0.04

0.00

0.00

10'0

S0

S1

S2

ES

S5

20

0.00

0.76

0.00

0.00

0.00

0.00

0.00

0.76

0.00

0.00

0.00

0.00

0.00

0.00

0.00

1.00

0.00

0.00

23

0.80

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.53

0.00

0.00

0.00

S

0.79

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.00

0.51

0.00

0.00

0.00

0.00

0.00

0.00

1.00

0.00

0.00

0.00

0.00

0.00

0.00

88'0

0.00

0.00

0.80

0.00

0.00

0.00

N

0.00

0.00

0.00

0.00

0.00

0.70

1lz

0.00

0.00

0.00

0.00

0.00

0.76

S0

S1

S2

ES

S5

1.0

0.66

0.03

00°0

100

0.07

0.08

100

100

10'0

0.61

0.18

0.01

2

0.25

100

00'0

0.01

0.03

0.23

0.8

0.26

0.02

000

000

0'0

0.28

0.17

10'0

00'0

00°0

0.02

0.22

0.6

0'0

0.00

10°0

0.04

0.02

0.13

90'0

0.03

0.00

0.01

0.13

0.02

0.4

0.03

10°0

1.00

0.01

0.02

0.05

0.03

1.00

00°0

0.01

0.04

0.04

3

0.02

0.12

00°0

0.07

0.72

0.07

0.2

0.14

0.02

00°0

0.01

0.10

0.62

0.09

00'0

0.02

0.03

0.03

0.07

0.0

S0

S1

S2

S3

S5

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 n s -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 [ 1 / n s , 1] , we re-normalize to [0 , 1] . Finally, we define InfoModularity (InfoM) as the average of this quantity over latents:

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 NMI instead of columns, and averaging over sources instead of latents, etc.:

We advocate for this terminology since previous names such as 'mutual information gap' [11] and 'mutual information ratio' [72] are ambiguous, and indeed the former of these works considered solely compactness and the latter solely modularity, with neither mentioning the distinction. We remark that when n z > n s (after pruning inactive latents), it is impossible to achieve both perfect modularity and perfect compactness. Of the two, modularity should be prioritized [58, 10] and indeed has been referred to as disentanglement itself [17].

4.2 Explicitness

Modularity and compactness are measured in terms of mutual information, so they are agnostic to how this information is encoded. Our third metric, explicitness, measures the extent to which the relationship between the sources and latents is simple (e.g., linear [41, 17, 59, 18]). Since previous explicitness metrics have been rather ad hoc, we propose a formalism using the framework of predictive 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 i given all latents z :

This requires estimating the predictive conditional V -entropy

and the marginal V -entropy of the source

where ∅ is an uninformative constant. The predictive conditional V -entropy measures how well a source, s i , 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 V 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 H V ( s i | ∅ ) in the same way, but substituting a universal constant for all inputs. We propose a simple normalization analogous to the one done for NMI :

which can be interpreted as the relative reduction in the V -entropy of a source achieved by knowing the latents, and is in [0 , 1] : for classification negative log likelihood is the cross-entropy, and for regression we leverage Propositions 1.3 and 1.5 from Xu et al. [75] to argue that NMI V ( z → s i ) = R 2 , the coefficient of determination. We can now compute explicitness as:

4.3 Summary and Comparison to Nonlinear DCI

We have derived three metrics for evaluating the modularity, explicitness, and compactness of a representation. Each metric has a straightforward information-theoretic interpretation and all share a range of [0 , 1] . We order them in decreasing importance and collectively refer to them as InfoMEC := ( InfoM , InfoE , InfoC ) .

Nonlinear DCI [17], a widely used three-pronged framework that measures similar disentanglement properties, suffers several practical drawbacks in comparison to InfoMEC from being defined in terms of relative counts of decision tree splits: determining this requires considering all latents jointly while fitting p ( s i | z ) . This results in a cumbersome computational footprint that is exacerbated by highly sensitive hyperparameters such as tree depth [17, 10], the tuning of which has even seen omission in prior work [47]. These drawbacks worsen with increased latent space dimensionality. In contrast, InfoMEC avoids these issues as it isolates InfoM and InfoC from the choice of predictive function class and only computes pairwise interactions between individual sources and latents. (InfoE also fits p ( s i | 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], MPI3D [20], Falcor3D [53], and Isaac3D [53]. Each consists of RGB image observations generated (near-)noiselessly from categorical or discretized numerical sources. Shapes3D is toyish, but the others are chosen for their difficulty [54, 20]. In particular, we use the complex shapes variant of MPI3D collected on a real world robotics apparatus. See Appendix C.1 for further dataset details. Aside from baseline AE and InfoGAN (specifically InfoWGAN-GP, a Wasserstein GAN [1] with gradient penalty [23]), we compare to β -VAE [25], β -TCVAE [11], and BioAE [72], the strongest methods from prior work that obey our problem assumptions and make design decisions mutually exclusive with latent quantization. We also compare to VQ-VAE with d = 64 and n v = 512 (Figure 2). We ablate weight decay, scalar codebooks, and dimension-specific codebooks from QLAE and weight decay from QLInfoWGAN-GP. We quantify modularity, explicitness, and compactness using both InfoMEC and nonlinear DCI. We qualitatively inspect representations via decoded latent interventions (Figure 4 and Appendix D).

Select experimental details. The choice of decoder architecture is known to impose inductive biases relevant for disentanglement [33, 43]. We use an expressive architecture (Appendix C.3) based on StyleGAN [33, 54] for all methods and datasets. We downsample the observations to 64 × 64 (if necessary). We follow prior work [47] in considering a statistical learning problem rather than a machine learning one: we train on the entire dataset then evaluate on 10 , 000 i.i.d. samples. We fix

the number of latents in all methods to twice the number of sources. For quantized-latent models, we fix n v = 10 discrete values per codebook. We tune one key regularization hyperparameter per method per dataset with a thorough sweep (Table 11, Appendix C.2). We use the best performing configurations over 2 seeds and rerun with 5 more seeds. Despite our modest list of methods and ablations, just the last stage took over 1000 GPU-hours.

Effect of latent quantization. Adding latent quantization to AE and InfoWGAN-GP results in consistent and dramatic increases in modularity and compactness under both InfoMEC and DCI evaluation (Table 1). For explicitness, the improvement is more significant under (random forest) I than under (linear) InfoE. Qualitatively, decoded latent interventions demonstrate that the QLAE latent space is highly interpretable, corroborating the gains in modularity (Figure 4 and Appendix D).

Comparison of QLAE with prior methods. QLAE significantly outperforms all prior methods on all four datasets in modularity under both InfoM and D (Table 1), with the exception of β -VAE and β -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 β -VAE and β -TCVAE contain a term that explicitly minimizes the total correlation (aka multiinformation), amongst the latent variables [11]. Since this essentially optimizes for compactness, we should expect compactness metrics to rank β -VAE and β -TCVAE ahead of QLAE; InfoC does so, but C does not. Corresponding metrics from InfoMEC and DCI have Spearman rank correlations of ρ = 0 . 80 , p < 1 × 10 -11 for modularity, ρ = 0 . 77 , p < 1 × 10 -9 for explicitness, and ρ = 0 . 59 , p < 1 × 10 -5 for compactness.

Ablations on latent space design and regularization. We observe that ablating dimension-specific codebooks, weight decay, and scalar codebooks from QLAE each causes a significant drop in InfoM and D (Table 2), verifying the importance of these design decisions. The effect of ablating weight decay from QLInfoWGAN-GP is less pronounced. We note that while VQ-VAE w/ weight decay performs somewhat closely to QLAE in terms of InfoMEC, this is only because we use the categorical codes (as opposed to the high-dimensional vector representation) for evaluation. In addition, the scalar codebook design of QLAE enables meaningful interpolation between discrete values, whereas this is not supported by vector quantization.

QLAE

BioAE

z

0

z

1

z

2

z

3

z

4

z

5

z

6

z

7

z

8

z

9

z

10

z

11

z

12

z

13

z

14

z

15

z

16

z

17

z

0

z

1

z

2

z

3

z

4

z

5

z

6

z

7

z

8

z

9

z

10

z

11

z

12

z

13

z

14

z

15

z

16

z

17

6 Related Work

Nonlinear ICA and disentangled representation learning. There is a long history of trying to build interpretable representations that separately represent the sources of variation in a dataset. This goes back to classic work on (linear) ICA [13, 29], and has been known in deep learning as disentanglement [4]. Without further assumptions, nonlinear ICA and its relatives in disentangled representation learning are provably underspecified [30, 47]. Approaches to resolve this indeterminacy include labeling a small number of datapoints [54] or showing pairs of datapoints in which only one or a few sources differ [63, 48]. More in line with our work are approaches that assume additional structure in the data generative process and disentangle by ensuring the representation reflects this structure [74]. Assumptions and methods include: factorized priors [25, 11, 36, 61], biologically inspired activity constraints [72], sparse source variation over time [65, 38], structurally sparse source to pixel influence [57, 31, 78, 52, 8], geometric assumptions on the source to image mapping [64, 77, 26, 22], sparse underlying causal graphs between sources [42], piecewise linearity [37], and hierarchical generation [43]. Many of these ideas are in principle compatible with latent quantization, and we leave discovery of fruitful combinations to future work.

Factorized latent spaces. VAEs that specify an isotropic Gaussian latent prior regularize the marginalized variational distribution (aka aggregate posterior) towards being a factorized distribution [11]. Roth et al. [61] bias latents to have pairwise factorized support via a Hausdorff set distance regularization. For linear ICA, Whittington et al. [72] prove that regularizing latents to be nonnegative and energy-minimizing results in them having factorized support. Differently from all of these works, latent quantization imbues a model with factorized structure by construction, instead of relying on the optimization of regularized objectives to manifest this structure. The favorable disentanglement that QLAE yields over β -TCVAE [11] and BioAE [72] suggests that this strategy is more effective.

Discrete representation learning. Oord et al. [55] first demonstrated the feasibility of discrete neural representation learning at scale, and their techniques have since been broadly applied, e.g., to videos [70, 76], audio [2, 15, 68, 6], and anomaly detection [50]. The following works design discrete representations similarly to how we do, though for purposes other than unsupervised disentanglement. Several works use one scalar codebook per latent dimension to achieve high efficiency in retrieval [3, 67, 32, 73]. Kobayashi et al. [39] disentangle normal and abnormal features in medical images into separate vector codebooks via pixel-space supervision. Liu et al. [46] and Träuble et al. [69] use multiple codebooks with separately parameterized key and value vectors and investigate the effect of discretization in systematic generalization and continual learning, respectively.

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-theoretic disentanglement metrics that rectify or sidestep key drawbacks in existing approaches.

While our results are promising, one concern might be that we have overfit our inductive bias to existing disentanglement benchmarks, in which, just like our model, the sources are discrete and the generative process is (near-)noiseless. Our experiments have already demonstrated the ability of latent quantization to represent sources that have more values (up to 40) than the per-dimension codebook size (fixed to 10) via allocating multiple latent dimensions. Future work should strive to construct disentanglement benchmarks that better reflect realistic conditions, e.g. continuous sources.

Beyond the intuitions and connections to related works we have presented, we do not understand why our method performs as well as it does. It may be fruitful to tackle this empirically, e.g. by probing how QLAE distributes data around its latent space, and how weight decay changes this. Achieving satisfactory understanding would enable the field to better position latent quantization within the ongoing body of work that aims to develop generalizable conditions for successful disentanglement.

Lastly, we hope this method, its future versions, and other methods the field develops are able to deliver on the original motivation for disentangled representation learning-to learn human-interpretable representations in complex, real-world situations, and to leverage the interpretability to empower human decision-making. This will require methods that can disentangle out-of-distribution data samples, that work for generic data types, and that can learn compositionally from sparse interactions with data. We suspect that latent quantization may have a role to play in these directions.

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], matplotlib [28], seaborn [71], and scikit-learn [56]. We also thank Evan Liu, Kaylee Burns, Karsten Roth, Anirudh Goyal, and Cian Eastwood for feedback on previous drafts.

KH was funded by a Sequoia Capital Stanford Graduate Fellowship. Part of WD's work on this project happened during a visit to Stanford funded by the Bogue Fellowship. JCRW was funded by a Henry Wellcome Post-doctoral Fellowship (222817/Z/21/Z). This work was also in part supported by the Stanford Institute for Human-Centered AI (HAI), NSF RI #2211258, Air Force Office of Scientific Research (AFOSR) YIP FA9550-23-1-0127, and ONR MURI N00014-22-1-2740.

References

Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning , 2017.

Alexei Baevski, Steffen Schneider, and Michael Auli. vq-wav2vec: self-supervised learning of discrete speech representations. In International Conference on Learning Representations , 2020.

Johannes Ballé, Valero Laparra, and Eero P. Simoncelli. End-to-end optimized image compression. In International Conference on Learning Representations , 2017.

Yoshua Bengio. Deep learning of representations: looking forward. In Statistical Language and Speech Processing , 2013.

Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432 , 2013.

Zalán Borsos, Raphaël Marinier, Damien Vincent, Eugene Kharitonov, Olivier Pietquin, Matt Sharifi, Olivier Teboul, David Grangier, Marco Tagliasacchi, and Neil Zeghidour. AudioLM: a language modeling approach to audio generation. arXiv preprint arXiv:2209.03143 , 2022.

James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax .

Simon Buchholz, Michel Besserve, and Bernhard Schölkopf. Function classes for identifiable nonlinear independent component analysis. In Advances in Neural Information Processing Systems , 2022.

Chris Burgess and Hyunjik Kim. 3D shapes dataset, 2018. URL https://github.com/ deepmind/3dshapes-dataset/ .

Marc-André Carbonneau, Julian Zaidi, Jonathan Boilard, and Ghyslain Gagnon. Measuring disentanglement: a review of metrics. IEEE Transactions on Neural Networks and Learning Systems , 2022.

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.

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.

Pierre Comon. Independent component analysis, a new concept? Signal processing , 36(3): 287-314, 1994.

Elliot Creager, David Madras, Joern-Henrik Jacobsen, Marissa Weis, Kevin Swersky, Toniann Pitassi, and Richard Zemel. Flexibly fair representation learning by disentanglement. In International Conference on Machine Learning , 2019.

Prafulla Dhariwal, Heewoo Jun, Christine Payne, Jong Wook Kim, Alec Radford, and Ilya Sutskever. Jukebox: a generative model for music. arXiv preprint arXiv:2005.00341 , 2020.

Mengnan Du, Ninghao Liu, and Xia Hu. Techniques for interpretable machine learning. Communications of the ACM , 63(1):68-77, 2019.

Cian Eastwood and Christopher KI Williams. A framework for the quantitative evaluation of disentangled representations. In International Conference on Learning Representations , 2018.

Cian Eastwood, Andrei Liviu Nicolicioiu, Julius Von Kügelgen, Armin Keki´ c, Frederik Träuble, Andrea Dittadi, and Bernhard Schölkopf. Dci-es: An extended disentanglement framework with connections to identifiability. In International Conference on Learning Representations , 2023.

Weihao Gao, Sewoong Oh, and Pramod Viswanath. Demystifying fixed k -nearest neighbor information estimators. IEEE Transactions on Information Theory , 64(8):5629-5661, 2018.

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.

Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. In Advances in Neural Information Processing Systems , 2014.

Luigi Gresele, Julius V on Kügelgen, Vincent Stimper, Bernhard Schölkopf, and Michel Besserve. Independent mechanism analysis, a new concept? In Advances in Neural Information Processing Systems , 2021.

Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C Courville. Improved training of Wasserstein GANs. In Advances in Neural Information Processing Systems , 2017.

Charles R Harris, K Jarrod Millman, Stéfan J Van Der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J Smith, et al. Array programming with NumPy. Nature , 585(7825):357-362, 2020.

Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. β -VAE: learning basic visual concepts with a constrained variational framework. In International Conference on Learning Representations , 2017.

Daniella Horan, Eitan Richardson, and Yair Weiss. When is unsupervised disentanglement possible? In Advances in Neural Information Processing Systems , 2021.

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.

John D Hunter. Matplotlib: a 2D graphics environment. Computing in Science & Engineering , 9(03):90-95, 2007.

Aapo Hyvärinen and Erkki Oja. Independent component analysis: algorithms and applications. Neural Networks , 13(4-5):411-430, 2000.

Aapo Hyvärinen and Petteri Pajunen. Nonlinear independent component analysis: existence and uniqueness results. Neural Networks , 12(3):429-439, 1999.

Li Jing, Jure Zbontar, et al. Implicit rank-minimizing autoencoder. In Advances in Neural Information Processing Systems , 2020.

Lukasz Kaiser, Samy Bengio, Aurko Roy, Ashish Vaswani, Niki Parmar, Jakob Uszkoreit, and Noam Shazeer. Fast decoding in sequence models using discrete latent variables. In International Conference on Machine Learning , 2018.

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.

Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational autoencoders and nonlinear ICA: a unifying framework. In International Conference on Artificial Intelligence and Statistics , 2020.

Patrick Kidger and Cristian Garcia. Equinox: neural networks in JAX via callable PyTrees and filtered transformations. Differentiable Programming Workshop at Neural Information Processing Systems , 2021.

Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In International Conference on Machine Learning , pages 2649-2658. PMLR, 2018.

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.

David A. Klindt, Lukas Schott, Yash Sharma, Ivan Ustyuzhaninov, Wieland Brendel, Matthias Bethge, and Dylan Paiton. Towards nonlinear disentanglement in natural data with temporal sparse coding. In International Conference on Learning Representations , 2021.

Kazuma Kobayashi, Ryuichiro Hataya, Yusuke Kurose, Mototaka Miyake, Masamichi Takahashi, 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.

Alexander Kraskov, Harald Stögbauer, and Peter Grassberger. Estimating mutual information. Physical Review E , 69(6):066138, 2004.

Abhishek Kumar, Prasanna Sattigeri, and Avinash Balakrishnan. Variational inference of disentangled latent concepts from unlabeled observations. In International Conference on Learning Representations , 2018.

Sebastien Lachapelle, Pau Rodriguez, Yash Sharma, Katie E Everett, Rémi Le Priol, Alexandre Lacoste, and Simon Lacoste-Julien. Disentanglement via mechanism sparsity regularization: a new principle for nonlinear ICA. In Conference on Causal Learning and Reasoning , 2022.

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.

Piyawat Lertvittayakumjorn and Francesca Toni. Explanation-based human debugging of NLP models: a survey. Transactions of the Association for Computational Linguistics , 9:1508-1528, 2021.

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.

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.

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.

Francesco Locatello, Ben Poole, Gunnar Rätsch, Bernhard Schölkopf, Olivier Bachem, and Michael Tschannen. Weakly-supervised disentanglement without compromises. In International Conference on Machine Learning , 2020.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations , 2019.

Sergio Naval Marimont and Giacomo Tarroni. Anomaly detection through latent space restoration using vector quantized variational autoencoders. In International Symposium on Biomedical Imaging , 2021.

Brent Mittelstadt. Auditing for transparency in content personalization systems. International Journal of Communication , 10:12, 2016.

Gemma Elyse Moran, Dhanya Sridhar, Yixin Wang, and David Blei. Identifiable deep generative models via sparse decoding. Transactions on Machine Learning Research , 2022.

Weili Nie. High resolution disentanglement datasets, 2019. URL https://github.com/ NVlabs/High-res-disentanglement-datasets .

Weili Nie, Tero Karras, Animesh Garg, Shoubhik Debnath, Anjul Patney, Ankit B Patel, and Anima Anandkumar. Semi-supervised StyleGAN for disentanglement learning. In International Conference on Machine Learning , 2020.

Aäron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learning. In Advances in Neural Information Processing Systems , 2017.

Fabian Pedregosa, Gaël Varoquaux, Alexandre Gramfort, Vincent Michel, Bertrand Thirion, Olivier Grisel, Mathieu Blondel, Peter Prettenhofer, Ron Weiss, Vincent Dubourg, Jake Vanderplas, Alexandre Passos, David Cournapeau, Matthieu Brucher, Matthieu Perrot, and Édouard Duchesnay. Scikit-learn: machine learning in Python. Journal of Machine Learning Research , 12(85):2825-2830, 2011.

William Peebles, John Peebles, Jun-Yan Zhu, Alexei Efros, and Antonio Torralba. The Hessian penalty: A weak prior for unsupervised disentanglement. In European Conference on Computer Vision , 2020.

Karl Ridgeway. A survey of inductive biases for factorial representation-learning. arXiv preprint arXiv:1612.05299 , 2016.

Karl Ridgeway and Michael C Mozer. Learning deep disentangled embeddings with the F-statistic loss. In Advances in Neural Information Processing Systems , 2018.

Brian C Ross. Mutual information between discrete and continuous data sets. PLOS ONE , 9(2): 1-5, 2014.

Karsten Roth, Mark Ibrahim, Zeynep Akata, Pascal Vincent, and Diane Bouchacourt. Disentanglement of correlated factors via hausdorff factorized support. In International Conference on Learning Representations , 2022.

Philipp Schmidt and Felix Biessmann. Quantifying interpretability and trust in machine learning systems. In AAAI Workshop on Network Interpretability for Deep Learning , 2019.

Rui Shu, Yining Chen, Abhishek Kumar, Stefano Ermon, and Ben Poole. Weakly supervised disentanglement with guarantees. In International Conference on Learning Representations , 2020.

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.

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.

Alexander Strehl and Joydeep Ghosh. Cluster ensembles-a knowledge reuse framework for combining multiple partitions. Journal of Machine Learning Research , 3(Dec):583-617, 2002.

Lucas Theis, Wenzhe Shi, Andrew Cunningham, and Ferenc Huszár. Lossy image compression with compressive autoencoders. In International Conference on Learning Representations , 2017.

Andros Tjandra, Sakriani Sakti, and Satoshi Nakamura. Transformer VQ-VAE for unsupervised unit discovery and speech synthesis: ZeroSpeech 2020 challenge. In INTERSPEECH , 2020.

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.

Jacob Walker, Ali Razavi, and Aäron van den Oord. Predicting video with VQVAE. arXiv preprint arXiv:2103.01950 , 2021.

Michael L Waskom. Seaborn: statistical data visualization. Journal of Open Source Software , 6 (60):3021, 2021.

James C R 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.

Hanwei Wu and Markus Flierl. Learning product codebooks using vector-quantized autoencoders for image retrieval. In Global Conference on Signal and Information Processing , 2019.

Quanhan Xi and Benjamin Bloem-Reddy. Indeterminacy in generative models: characterization and strong identifiability. In International Conference on Artificial Intelligence and Statistics , 2023.

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.

Wilson Yan, Yunzhi Zhang, Pieter Abbeel, and Aravind Srinivas. VideoGPT: video generation using VQ-VAE and transformers. arXiv preprint arXiv:2104.10157 , 2021.

Xiaojiang Yang, Yi Wang, Jiacheng Sun, Xing Zhang, Shifeng Zhang, Zhenguo Li, and Junchi Yan. Nonlinear ICA using volume-preserving transformations. In International Conference on Learning Representations , 2022.

Yujia Zheng, Ignavier Ng, and Kun Zhang. On the identifiability of nonlinear ICA: sparsity and beyond. In Advances in Neural Information Processing Systems , 2022.

A Quantized-Latent Models

This section contains pseudocode for latent quantization and for training QLAE and QLInfoWGAN-GP.

B Disentanglement Metrics Vignettes

0

25

50

75

100

0.0

0.2

0.4

0.6

0.8

1.0

BioAE, Shapes3D

InfoM

InfoC

0

25

50

75

100

number of bins

-TCVAE, Falcor3D

0

25

50

75

100

-VAE, Isaac3D

0

25

50

75

100

0.0

0.2

0.4

0.6

0.8

1.0

BioAE, Shapes3D

InfoM

InfoC

0

25

50

75

100

number of neighbors

-TCVAE, Falcor3D

0

25

50

75

100

-VAE, Isaac3D

0

5

10

15

20

0.0

0.2

0.4

0.6

0.8

1.0

BioAE, Shapes3D

D

C

I

0

5

10

15

20

maximum tree depth

-TCVAE, Falcor3D

0

5

10

15

20

-VAE, Isaac3D

C Experiment Details

This section contains details on the experiments conducted in this work.

C.1 Datasets

！

C.2 Hyperparameters

This section specifies fixed and tuned hyperparameters for all methods considered.

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 (kernel size 3 , stride 1 ) and one downsampling convolutional layer (kernel size 4 , stride 2 ) at a consistent width (number of channels). Each convolution operation is followed by a leaky ReLU (slope 0 . 3 ), then instance normalization. There are four such blocks with widths 32 , 64 , 128 , and 256 . The 256 × 4 × 4 output is then flattened. Two dense layers each of width 256 with ReLU activation (and no normalization) follow. The final operation is an affine projection to the latent layer.

Decoder. We use a decoder architecture based on StyleGAN [33]. The latent code z of shape n z passes through two dense layers each of width 256 with ReLU activation (and no normalization); call the output of this w . We directly parameterize a starting input feature map of shape 256 × 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 a leaky ReLU (slope 0 . 3 ) and then an adaptive instance normalization (AdaIN): w undergoes an affine projection to a scale and bias scalar for each channel of the output feature map, and the output feature map is 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-preserving style-decoding layers (kernel size 3 , stride 1 ) and one upsampling style-decoding layer (kernel size 4 , stride 2 ) at a consistent width. There are four such blocks with widths 256 , 128 , 64 , and 32 . The final operation is a 1 × 1 convolution of width 3 to yield an output of shape 3 × 64 × 64 .

C.4 Negative Results

We tried a number of additional modifications to our basic methods, QLAE and QLInfoWGAN-GP, beyond the ablations presented in the main paper. Here is a list of those that only marginally helped, didn't help, or made things worse:

Keeping the original VQ-VAE [55] hyperparameter settings for λ quantize ( 1 ) and λ commit ( 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 λ quantize and/or λ commit from 0 up to a specified value.

Using ℓ 1 norm parameter vector regularization, in addition to or without weight decay.

(QLInfoWGAN-GP-specific) Updating V with the encoder instead of with the decoder.

C.5 Nonlinear DCI Evaluation

For each p ( s i | z ) learning problem, we train a random forest classifier with 100 trees and the information gain splitting criterion. We use a 0 . 9 / 0 . 1 train/test split of the evaluation sample. We tune the maximum tree depth hyperparameter on held-out accuracy (see Figure 7 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 I.

D Qualitative Results

QLAE

-TCVAE

floor color

0.00

0.00

0.00

0.80

0.00

0.79

0.00

00'0

0.00

0.00

00'0

0.00

floor color

0.66

0.01

0.25

0.26

0.17

0.05

0.06

0.03

0.03

10'0

0.14

0.09

object color

0.76

0.76

0.00

0.00

000

0.00

0.00

00'0

00°0

0.00

0.00

0.00

object color

0.03

0.01

10°0

0.02

10'0

00'0

0.03

10'0

1.00

0.12

0.02

00'0

camera orientation

0.00

0.00

0.00

0.00

0.53

0.00

0.00

0.00

0.00

0.80

0.00

0.00

camera orientation

0.00

0.01

0.00

0.01

0.00

10'0

0.00

1.00

0.00

00°0

0.00

0.02

object scale

0.00

0.00

1.00

0.00

0.00

0.00

0.51

0.00

0.00

0.00

0.00

0.00

object scale

0.01

0.61

0.01

00°0

0.01

0.04

0.01

100

0.02

0.07

0.01

0.03

object shape

0.00

0.00

0.00

0.00

00'0

0.00

0.00

1.00

0.00

0.00

00'0

0.00

object shape

0.07

0.18

0.03

0'0

0.02

0.02

0.13

0.02

0'0

0.72

0.10

0.03

wall color

000

0.00

00'0

0.00

0.00

00'0

0.00

88'0

0.00

0.70

0.76

wall color

800

0.02

0.23

0.28

0.22

0.13

20'0

0.05

0.04

0.07

0.62

0.07

Zo

23

28

Z9

Z11

Zo

L2

43

L7

Z8

Z9

Ilz

QLAE

-TCVAE

object color

0.000.01

0.98

0.000.000.000.000.000.010.000.010.000.010.00

object color

0.020.02

0.72

0.000.020.100.020.060.010.030.000.010.000.00

object shape

0.000.010.030.000.000.000.020.01

0.46

0.00

0.400.00

0.35

0.00

object shape

0.39

0.080.060.080.010.010.040.030.050.030.110.040.060.14

object size

0.000.010.020.000.000.000.020.00

0.25

0.000.050.00

0.51

0.00

object size

0.160.000.020.030.010.020.020.010.050.010.040.000.010.06

camera height

0.00

0.73

0.020.000.000.010.05

0.60

0.010.000.010.000.000.00

camera height

0.050.040.060.010.010.010.03

0.33

0.00

0.89

0.01

0.17

0.020.00

background color

0.000.010.000.000.00

0.00

1.00

0.000.000.000.000.000.000.00

background color

S00200200100

0.40

0000

0.76

0.01

0.020.000.040.01

0.00

robot x

0.060.010.010.030.06

0.03

0.01

0.04

100

0.070.020.020.01

0'0

robot x

0.000.010.000.010.010.000.010.000.060.010.050.010.080.02

roboty

0.070.030.010.060.05

0.11

0.02

600

500

00

0'0

0.15

0.02

0.05

robot y

0°0

0.19

Z12

Z13

Z9Z10Z11Z12Z13

QLAE

-VAE

lighting intensity

0.020.010.020.050.11

0.76

0.110.010.000.020.020.070.030.07

lighting intensity

0.73

0.030.020.110.010.020.000.020.030.030.020.120.060.05

lighting x

0.000.010.040.050.050.010.040.010.010.010.00

0.44

0.010.11

lighting x

0.010.000.000.020.020.030.000.020.05

0.52

100100100

lighting y

0.010.000.080.050.030.010.060.000.000.000.000.040.00

0.19

lighting y

0.030.020.010.000.010.010.010.00

0.150.17

0.01

0.01

0.03

lighting z

0.010.01

0.370.460.54

0.020.230.030.050.000.010.020.010.04

lighting z

0.120.100.080.000.100.030.020.03

0.62

0.050.050.030.030.12

camera x

0.030.020.000.000.000.000.00

0.700.74

0.02

0.000.010.00

camera x

0.000.050.010.100.040.02

0.18

0.050.000.010.110.090.14

1.00

camera y

0.010.020.000.000.000.010.000.000.01

0.78

0.000.00

0.84

0.00

cameray

0.000.09 0.090.07

1.00

0.17

0.080.080.010.000.060.030.050.04

camera Z

0.750.80

0.050.000.000.010.010.020.010.010.010.000.010.00

camera Z

0.010.050.150.060.130.060.080.15

1.00

0.040.050.07

ZoZ1Z2Z3Z4Z5

Z6

Z7Z8Z9Z10 Z11 Z12 Z13

Zo

Z1Z2

Z3

Z4

Z6

Z7

Z8

Z9Z10 Z11 Z12 Z13

QLAE

BioAE

上

上

上

object shape

0.00 0.03 0.00

0.59 0.31 0.00 0.02 0.05 0.02 0.05 0.00 0.05 0.01 0.02 0.11 0.00 0.01 0.05

object shape

0.00 0.00 0.02 0.01 0.03 0.01 0.03 0.00 0.00 0.00 0.01 0.03 0.04 0.05 0.03 0.00 0.00 0.03

robot x

0.00 0.03

0.67 0.03 0.03 0.02

0.540.01 0.03 0.00 0.03 0.03 0.01 0.04 0.01 0.01 0.03 0.01

robot x

0.01 0.02 0.45 0.11 0.05 0.02 0.43 0.03 0.01 0.17 0.01 0.09 0.02 0.09 0.08 0.03 0.01 0.03

roboty

0.00

0.83

0.06 0.02 0.00 0.00 0.03 0.00

0.67

0.00 0.040.470.00

0.67 0.01 0.00 0.01 0.00

roboty

0.01 0.00 0.04 0.28 0.50 0.01 0.07 0.03 0.01 0.24 0.02 0.03 0.01 0.06 0.03 0.02 0.01 0.01

camera height

0.01 0.01 0.01 0.01 0.01 0.02 0.01 0.00 0.03 0.00

0.75

0.01 0.01 0.03 0.00 0.02

.75

00'0

camera height

0.04 0.04 0.12 0.01 0.02 0.05 0.02 0.14 0.05 0.05 0.03

0.69

0.000.980.77

10°0500600

object scale

0.01 0.00 0.01 0.020.390.00 0.05 0.09 0.000.280.01 0.01 0.00 0.000.97

0.00 0.01 0.11

object scale

0.00 0.00 0.01 0.01 0.02 0.00 0.01 0.00 0.00 0.01 0.00 0.14 0.08 0.26 0.15 0.00 0.01 0.07

lighting intensity

1.00

0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.02 0.00 0.03 0.00 0.00 0.020.00 0.00

lighting intensity

0.480.20 0.01 0.02 0.01 0.27 0.00 0.19 0.79

0.02 0.06 0.00 0.00 0.01 0.01 0.11 0.31

000

lighting direction

0.76

0.01 0.00 0.01 0.01 0.00

lighting direction

0.050.710.00 0.03 0.000

0.81

0.00 0.00 0.31 0.00 0.25 0.00 0.00 0.01 0.00 0.17

0.69

0.00

object color

0.00 0.00 0.01 0.01 0.01 0.00 0.01

0.97

0.00 0.34 0.03 0.01 0.00 0.00 0.01 0.01 0.00

0.97

object color

0.01 0.01 0.02 0.01 0.01 0.01 0.00 0.01 0.01 0.01 0.00 0.020.42 0.02 0.04 0.01 0.01

0.39

wall color

00°000°000000010000000°000020000000000°000°0000100

1.00

0.00 0.00

wall color

0.08 0.02 0.01 0.01 0.00 0.01 0.00 0.00 0.02 0.00 0.01 0.01 0.00 0.00 0.001.000

0.010.00

Zo Z1 Z2 Z3 Z4 Z5 Z6 Z7 Z8 Z9 Z10 Z11 Z12 Z13 Z14 Z15 Z16 Z17

Zo Z1 Z2 Z3 Z4 Z5 Z6 Z7 Z8 Z9 Z10 Z11 Z12 Z13 Z14 Z15 Z16 Z17

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-perfect data reconstruction, whereas InfoGANs are generally more lossy.

Peer reviews:
comments: 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.

STRENGTHS:
The paper proposed a new metric for disentanglement.

WEAKNESSES:
1) The manuscript could benefit from improvements as some parts and notations are rather confusing.

2) The proposed method appears to be very similar to VQ-VAE.

3) It is uncertain if the underlying assumption that real-world data is generated in a compositional manner holds.

4) 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.

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?

2) The new interpretation given in equation 3 for infoGAN does not appear convincing.

3) 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-03

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nIn 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.\n\nWEAKNESSES:\nOn 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).\n\nQUESTIONS:\nFor 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.\n\nLIMITATIONS:\nThere 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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\nWEAKNESSES:\nThe 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\n\nQUESTIONS:\nDoes 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\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n— 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.\n\nWEAKNESSES:\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\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.\n\nQUESTIONS:\nPlease see questions in 'Weakness' section.\n\nLIMITATIONS:\nI did not realize potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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;\n\nSTRENGTHS:\n- 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;\n\nWEAKNESSES:\n- 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;\n\nQUESTIONS:\nQuestions: \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;\n\nLIMITATIONS:\nI 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;", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-25"}