SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Kyle Hsu† Will Dorrell‡ James C. R. Whittington†§ Jiajun Wu† Chelsea Finn†
†Stanford University ‡University College London ‡Oxford University
{kylehsu, jiajunwu, cbfinn}@cs.stanford.edu {dorrellwec, jcrwhittington}@gmail.com
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.
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)—continuous codes (dots) are mapped (via arrows) to their nearest discrete codes (black circles), each of which is 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 of 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 a comparatively local effect.
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 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 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 data-reconstructing (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.<sup>1</sup>
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:
$$
p (\mathbf {s}) = \prod_ {i = 1} ^ {n _ {s}} p \left(\mathbf {s} _ {i}\right), \quad \mathbf {x} = g (\mathbf {s}), \tag {1}
$$
$^{1}$ Code for models and InfoMEC metrics: https://github.com/kylehkhsu/latentquantization.
where  $\mathbf{s} = (\mathbf{s}_1, \dots, \mathbf{s}_{n_s})$  comprise the  $n_s$  mutually independent source variables (sources);  $\mathbf{x}$  is the observed data variable; and  $g: S \to \mathcal{X}$  is the nonlinear data-generating function. The nonlinear ICA problem is to recover the underlying sources given a dataset  $\mathcal{D}$  of samples from this model. Specifically, the solution should include an approximate inverting function  $\hat{g}^{-1}: \mathcal{X} \to \mathcal{Z}$  such that, assuming latent variables (latents)  $\mathbf{z} = (\mathbf{z}_1, \dots, \mathbf{z}_{n_s})$  are matched correctly to the sources, each source is perfectly determined by its corresponding latent. A typical technical phrasing is that  $\hat{g}^{-1} \circ g$  should be the composition of a permutation and a dimension-wise invertible function.
As stated, this problem is nonidentifiable (or underspecified). Given  $\mathcal{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  $\hat{g}^{-1}:\mathcal{X}\to \mathcal{Z}$  and decoder  $\hat{g}:\mathcal{Z}\rightarrow \mathcal{X}$  solely using the data  $\{x\}$  such that (i) the InfoMEC as estimated from samples  $\{(s,z = \hat{g}^{-1}\circ g(s))\}$  from the joint source-latent distribution is high, while (ii) maintaining an acceptable level of reconstruction error between  $x$  and  $\hat{g}\circ \hat{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  $\hat{g}^{-1}$  and decoder  $\hat{g}$ . An autoencoder takes a datapoint  $x\in \mathcal{X}$  as input and produces a reconstruction  $\hat{g}\circ \hat{g}^{-1}(x)\in \mathcal{X}$  that is optimized to match the input:
$$
\mathcal {L} _ {\text {r e c o n s t r u c t d a t a}} \left(\hat {g} ^ {- 1}, \hat {g}; \mathcal {D}\right) := \mathbb {E} _ {x \sim \mathcal {D}} \left[ - \log p (x \mid \hat {g} \circ \hat {g} ^ {- 1} (x)) \right]. \tag {2}
$$
An InfoGAN instead takes a latent code  $z \in \mathcal{Z}$  as input. The decoder (aka generator) maps  $z$  to the data space, and from this the encoder produces a reconstruction of the latent:
$$
\mathcal {L} _ {\text {r e c o n s t r u c t l a t e n t}} (\hat {g} ^ {- 1}, \hat {g}) := \mathbb {E} _ {z \sim p (\mathbf {z})} \left[ - \log p (z \mid \hat {g} ^ {- 1} \circ \hat {g} (z)) \right]. \tag {3}
$$
Unlike the data reconstruction loss, this is clearly insufficient for learning as the dataset  $\mathcal{D}$  isn't even used. InfoGAN can be thought of as grounding latent reconstruction by making the marginal distribution of generated datapoints,  $\hat{g} (\mathbf{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
Figure 2: Two ways to quantize into one of  $n_v^{n_z}$  ( $5^2$ ) discrete codes. (a) Vector quantization [55] splits a continuous representation of size  $n_z d$  (4) into  $n_z$  parts (black dots), each of which is quantized to the nearest of a global codebook of  $n_v$  vectors of size  $d$  (turquoise circles). (b) Latent quantization specifies  $d = 1$  and quantizes a continuous representation of size  $n_z$  (black dot) onto a regular grid parameterized by dimension-specific codebooks (orange and turquoise ticks), each of size  $n_v$ . Unlike vector quantization, latent quantization ties the decoder input space to the discrete code space.
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 \times \dots \times 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  $\hat{g}^{-1}: \mathcal{X} \to \mathbb{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:
$$
z _ {j} = \underset {v _ {j k} \in V _ {j}} {\arg \min } | \hat {g} ^ {- 1} (x) _ {j} - v _ {j k} |, j = 1, \dots , n _ {z}. \tag {4}
$$
We represent  $Z$  by its constituent discrete values, concretely as a learnable two-dimensional array  $V \in \mathbb{R}^{n_z \times 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:
$$
\mathcal {L} _ {\text {q u a n t i z e}} = \left\| \operatorname {S t o p G r a d i e n t} \left(\hat {g} ^ {- 1} (x)\right) - z \right\| _ {2} ^ {2}, \quad \mathcal {L} _ {\text {c o m m i t}} = \left\| \hat {g} ^ {- 1} (x) - \operatorname {S t o p G r a d i e n t} (z) \right\| _ {2} ^ {2}. \tag {5}
$$
The straight-through gradient estimator facilitates the flow of gradients through the nondifferentiable quantization step.  $\mathcal{L}_{\mathrm{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.
Algorithm 1 Latent quantization and computation of codebook losses.
1: function LatentQuantization(matpoint  $x$ , encoder  $\hat{g}^{-1}$ , discrete value array V)  
2:  $z_{c} \gets \hat{g}^{-1}(x)$   
3:  $z \gets \arg \min_{v \in Z} \| z_{c} - v \|_{1}, v_{j} \in V_{j}, \|_{j=1}^{n_{z}} V_{j} = V$   
4:  $\mathcal{L}_{\text{quantize}} \gets \|\text{StopGradient}(z_{c}) - z\|_{2}^{2}$   
5:  $\mathcal{L}_{\text{commit}} \gets \|z_{c} - \text{StopGradient}(z)\|_{2}^{2}$   
6:  $z \gets z_{c} + \text{StopGradient}(z - z_{c})$   
7: return  $z, \mathcal{L}_{\text{quantize}}, \mathcal{L}_{\text{commit}}$
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,
$$
I \left(\mathbf {s} _ {i}; \mathbf {z} _ {j}\right) := D _ {\mathrm {K L}} \left(p \left(\mathbf {s} _ {i}, \mathbf {z} _ {j}\right) \| p \left(\mathbf {s} _ {i}\right) p \left(\mathbf {z} _ {j}\right)\right), \tag {6}
$$
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(\mathbf{s}_i; \mathbf{z}_j) = H(\mathbf{s}_i) - H(\mathbf{s}_i \mid \mathbf{z}_j)$  and the nonnegativity of entropy imply  $I(\mathbf{s}_i; \mathbf{z}_j) \leq H(\mathbf{s}_i)$  for discrete sources. Following [11], we define a normalized mutual information as
$$
\operatorname {N M I} \left(\mathbf {s} _ {i}, \mathbf {z} _ {j}\right) := \frac {I \left(\mathbf {s} _ {i} ; \mathbf {z} _ {j}\right)}{H \left(\mathbf {s} _ {i}\right)}. \tag {7}
$$
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  $\mathrm{NMI}(\mathbf{s}_i,\mathbf{z}_j)$  into a 2-dimensional array  $\mathrm{NMI}\in [0,1]^{n_s\times 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  $\mathrm{NMI}^\top$ .
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
Figure 3: Visualization of  $\mathrm{NMI}^{\top}$  for Shapes3D models. Inactive latents (red font) are removed from InfoM and InfoC computation. The low values in the  $\mathrm{NMI}^{\top}$  of AE indicate that individual latents are 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  $\beta$ -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  $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  $\left[ \frac{1}{n_s}, 1 \right]$ , we re-normalize to  $[0, 1]$ . Finally, we define InfoModularity (InfoM) as the average of this quantity over latents:
$$
\operatorname {I n f o M} := \left(\frac {1}{n _ {z}} \sum_ {j = 1} ^ {n _ {z}} \frac {\max  _ {i} \mathrm {N M I} _ {i j}}{\sum_ {i = 1} ^ {n _ {s}} \mathrm {N M I} _ {i j}} - \frac {1}{n _ {s}}\right) / \left(1 - \frac {1}{n _ {s}}\right). \tag {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 NMI instead of columns, and averaging over sources instead of latents, etc.:
$$
\operatorname {I n f o C} := \left(\frac {1}{n _ {s}} \sum_ {i = 1} ^ {n _ {s}} \frac {\max  _ {j} \mathrm {N M I} _ {i j}}{\sum_ {j = 1} ^ {n _ {z}} \mathrm {N M I} _ {i j}} - \frac {1}{n _ {z}}\right) / \left(1 - \frac {1}{n _ {z}}\right). \tag {9}
$$