SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
p \left(\boldsymbol {z} \mid \boldsymbol {\theta} _ {d e c}\right) = p \left(\boldsymbol {z} _ {1}\right) \cdot \prod_ {\ell = 2} ^ {L} p \left(\boldsymbol {z} _ {\ell} \mid \boldsymbol {z} _ {<   \ell}; \boldsymbol {\theta} _ {d e c}\right), \tag {7}
$$
where the conditionals are represented by factorial Normal distributions. For the first latent group we have  $p(z_{1}) = \mathcal{N}(\mathbf{0},\mathbf{I})$ , similar to vanilla VAEs. For every other latent group, we have:
$$
p \left(\boldsymbol {z} _ {\ell} \mid \boldsymbol {z} _ {<   \ell}; \boldsymbol {\theta} _ {d e c}\right) = \mathcal {N} \left(\boldsymbol {\mu} \left(\boldsymbol {z} _ {<   \ell}; \boldsymbol {\theta} _ {d e c}\right), \boldsymbol {\sigma} ^ {2} \left(\boldsymbol {z} _ {<   \ell}; \boldsymbol {\theta} _ {d e c}\right)\right), \tag {8}
$$
where  $\mu (z_{<\ell};\theta_{dec})$  and  $\sigma^2 (z_{<\ell};\theta_{dec})$  are outputted from the decoder sampler layers. Similarly, the approximate posterior in the cNVAE is represented by:
$$
q (\boldsymbol {z} | \boldsymbol {x}; \boldsymbol {\theta} _ {e n c}) = \prod_ {\ell = 1} ^ {L} q \left(\boldsymbol {z} _ {\ell} \mid \boldsymbol {z} _ {<   \ell}, \boldsymbol {x}; \boldsymbol {\theta} _ {e n c}\right). \tag {9}
$$
We adopt a Gaussian parameterization for each conditional in the approximate posterior:
$$
q \left(\boldsymbol {z} _ {\ell} \mid \boldsymbol {z} _ {<   \ell}, \boldsymbol {x}; \boldsymbol {\theta} _ {e n c}\right) = \mathcal {N} \left(\boldsymbol {\mu} \left(\boldsymbol {z} _ {<   \ell}, x; \boldsymbol {\theta} _ {e n c}\right), \boldsymbol {\sigma} ^ {2} \left(\boldsymbol {z} _ {<   \ell}, \boldsymbol {x}; \boldsymbol {\theta} _ {e n c}\right)\right), \tag {10}
$$
where  $\pmb{\mu}(\pmb{z}_{<\ell}, x; \pmb{\theta}_{enc})$  and  $\sigma^2(\pmb{z}_{<\ell}, \pmb{x}; \pmb{\theta}_{enc})$  are outputted from the encoder sampler layers (Fig. 2; grey trapezoids). We are now in a position to explicitly write down the KL term from Equation 4 for the cNVAE:
$$
\operatorname {K L} \text {t e r m} = \mathcal {D} _ {\mathrm {K L}} \left[ q \left(\boldsymbol {z} _ {1} \mid \boldsymbol {x}, \boldsymbol {\theta} _ {\text {e n c}}\right) \| p \left(\boldsymbol {z} _ {1}\right) \right] + \sum_ {\ell = 2} ^ {L} \mathbb {E} _ {q \left(\boldsymbol {z} _ {<   \ell} \mid \boldsymbol {x}, \boldsymbol {\theta} _ {\text {e n c}}\right)} \left[ \mathrm {K L} _ {\ell} \left(\boldsymbol {\theta} _ {\text {e n c}}, \boldsymbol {\theta} _ {\text {d e c}}\right) \right], \tag {11}
$$
where  $\mathrm{KL}_{\ell}$  refers to the local KL term for group  $\ell$  and is given by:
$$
\mathrm {K L} _ {\ell} \left(\boldsymbol {\theta} _ {\text {e n c}}, \boldsymbol {\theta} _ {\text {d e c}}\right) := \mathcal {D} _ {\mathrm {K L}} \left[ q \left(\boldsymbol {z} _ {\ell} \mid \boldsymbol {z} _ {<   \ell}, \boldsymbol {x}; \boldsymbol {\theta} _ {\text {e n c}}\right) \| p \left(\boldsymbol {z} _ {\ell} \mid \boldsymbol {z} _ {<   \ell}, \boldsymbol {\theta} _ {\text {d e c}}\right) \right], \tag {12}
$$
and the approximate posterior up to the  $(\ell -1)^{th}$  group is defined as:
$$
q \left(\boldsymbol {z} _ {<   \ell} \mid \boldsymbol {x}; \boldsymbol {\theta} _ {e n c}\right) := \prod_ {i = 1} ^ {\ell - 1} q \left(\boldsymbol {z} _ {i} \mid \boldsymbol {z} _ {<   i}, \boldsymbol {x}; \boldsymbol {\theta} _ {e n c}\right). \tag {13}
$$
This concludes our derivation of the KL terms shown in Table 2.
9.2 Key architectural differences between NVAE and cNVAE
Our main contribution to architecture design lies in the modification of the NVAE "pager" layer and the introduction of the "expand" layer, which are simple deconvolution (Fig. 2). In the NVAE, the purpose of a sampler layer is to map encoder features to mean and variance vectors  $\mu(x)$  and  $\sigma^2(x)$ , which are then used to construct approximate posterior distributions  $q(z|x) = \mathcal{N}(\mu(x), \sigma^2(x))$ . During the inference forward pass, downsampling operations will progressively reduce the spatial scale of representations along the encoder pathway, such that processed stimuli will have different spatial dimensions at different stages of inference. In the NVAE, all sampler layers are convolutional with a fixed kernel size of  $3 \times 3$  and padding of 1. Thus, the application of a sampler layer does not alter the spatial dimension of hidden encoder features, resulting in a convolutional latent space.
For example, at level  $\ell$  of the hierarchy,  $\pmb{\mu}_{\ell}(\pmb{x})$  and  $\sigma_{\ell}^{2}(\pmb{x})$  will have shapes like  $d\times s_{\ell}\times s_{\ell}$ , where  $d$  is number of latent variables per latent group, and  $s_\ell$  is the spatial scale of the processed encoder features at level  $\ell$ . As a result, NVAE latents themselves end up being convolutional in nature, with shapes like  $d\times s_{\ell}\times s_{\ell}$ . This design choice leads to the massive expansion of NVAE latent dimensionality, which is often orders of magnitude larger than the input dimensionality.
To address this, we modified NVAE sampler layers in a manner that would integrate over spatial information before the sampling step, achieving a compressed, non-convolutional, and more abstract latent space. Specifically, we designed the cNVAE sampler layers such that their kernel sizes are selected adaptively to match the spatial scale of their input feature. For example, if latent group  $\ell$  operates at the scale of  $s_{\ell}$ , then the sampler layer will have a kernel size of  $s_{\ell} \times s_{\ell}$ , effectively integrating over space before the sampling step. This results in latents with shape  $d \times 1 \times 1$ .
Finally, combining the latents with decoder features requires projecting them back into convolutional space. We achieve this by introducing "expand" modules (Fig. 2; yellow trapezoids) which are simple deconvolution layers with a kernel size of  $s_{\ell} \times s_{\ell}$ . Thus, processing a  $d \times 1 \times 1$  latent by an expand layer results in a  $d \times s_{\ell} \times s_{\ell}$  output, allowing them to be combined with the decoder features (concatenation or addition; blue "+" in Fig. 2).
cNVAE pros and cons. Our approach resulted in a much smaller latent space compared to NVAE, which allowed us to examine the function of individual latents, decoupled from their spatial influence. However, one drawback of our approach is that it makes the architecture input-size-dependent. Scaling up the input increases the parameter count of cNVAE. This did not pose serious problems in the present work, as we generated ROFL datasets in a relatively small spatial scale of  $17 \times 17$ . However, this dependence on input dimensionality could cause problems when larger datasets are considered. To mitigate this, one possibility is to use depthwise separable convolutions for sampler and expand layers.
Other details. We used the same regular residual cells for both the encoder and the decoder (" $r$ " in Fig. 2). We did not experience memory issues because of the relatively small scale of our experiments. However, rendering the input data at a larger spatial dimension would lead to larger memory usage, which might require using depthwise separable convolutions as in Figure 3b in Vahdat and Kautz [52]. Similar to NVAE, we found that the inclusion of squeeze & excitation [126] helped. Contrary to the suggestion of NVAE [52] and others [72], batch norm [127] destabilized training in our case. Instead, we found that weight normalization [128] and swish activation [129, 130] were both instrumental in training. We used weight normalization without data-dependent initialization. Our observations also deviate from those of Vahdat and Kautz [52] with regard to spectral regularization [131]. We achieved stable training and good results without the need for spectral regularization. This is perhaps because our dataset is synthetic and has a much smaller dimensionality compared to real-world datasets such as faces considered in the NVAE. Finally, we also found that residual Normal parameterization of the approximate posterior was helpful.
Model hyperparameters. The latent space of cNVAE contains a hierarchy of latent groups that are sampled sequentially: "top" latents are sampled first, all the way down to "bottom" latents that are closest to the stimulus space (Fig. 2). Child [71] demonstrated that model depth, in this statistical sense, was responsible for the success of hierarchical VAEs. Here, we increased the depth of cNVAE architecture while ensuring that training remained stable. Our final model architecture had a total of 21 hierarchical latent groups: 3 groups operating at the spatial scale of  $2 \times 2$ , 6 groups at the scale of  $4 \times 4$ , and 12 groups at the scale of  $8 \times 8$  (see Figs. 2 and 3). Each latent group has 20 latents, which yields a  $21 \times 20 = 420$  dimensional latent space. See below for more details. We observed that various aspects of model performance such as reconstruction loss and DCI scores increased with depth, corroborating previous work [71]. We set the initial number of channels to 32 in the first layer of the encoder and doubled the channel size every time we downsample the features spatially. This resulted in a width of 256 in the final layer. Please refer to our code for more information.
9.3 cNVAE vs. NVAE latent dimensionality
As an illustrative example, here we count the number of latent variables in cNVAE and compare them to an otherwise identical but non-compressed NVAE. The total number of latent variables for a hierarchical VAE is given by:
$$
D = \sum_ {\ell = 1} ^ {L} d _ {\ell}, \tag {14}
$$
Table 4: The hierarchical models (cNVAE, cNAE) process information at various spatial scales. These models have a total of 21 latent groups, each containing  $d = 20$  latent variables. For hierarchical models, we concatenate the latent representations across all levels of the hierarchy, resulting in a single latent vector with dimensionality  $D = 420$ . In contrast, non-hierarchical models (VAE, AE) do not possess such a multi-scale organization. Instead, their non-hierarchical latent spaces contain a single latent group, operating at a single spatial scale. This single latent group contains the entire set of  $d = D = 420$  latent variables. See Table 2 and Fig. 2.
Model	Number of latent groups	Number of latent
variables per group	Latent space
dimensionality
2 × 2	4 × 4	8 × 8	total
cNVAE	3	6	12	21	d = 20	D = 21 × 20 = 420
VAE	1	-	-	1	d = 420	D = 1 × 420 = 420
cNAE	3	6	12	21	d = 20	D = 21 × 20 = 420
AE	1	-	-	1	d = 420	D = 1 × 420 = 420
where  $d_{\ell}$  is the number of latents in level  $\ell$ , and  $L$  is the total number of hierarchical levels. As described above,  $d_{\ell} = d \cdot s_{\ell}^{2}$  for NVAE, where  $d$  is the number of latents per level and  $s_{\ell}$  is the spatial scale of level  $\ell$ . In contrast, this number is a constant for the cNVAE:  $d_{\ell} = d$ . In this paper, we set  $d = 20$  and had the following configuration for the cNVAE:
- "top" latents: 3 groups operating at the spatial scale of  $2 \times 2$
- "mid" latents: 6 groups operating at the spatial scale of  $4 \times 4$
- "bottom" latents: 12 groups operating at the spatial scale of  $8 \times 8$
This resulted in an overall latent dimensionality of  $D = (3 + 6 + 12) \times d = 21 \times 20 = 420$  for the cNVAE. In contrast, an NVAE with a similar number of hierarchical levels would have a dimensionality of  $D = (3 \cdot 2^2 + 6 \cdot 4^2 + 12 \cdot 8^2) \times d = 876 \times 20 = 17,520$ .
Note that in Fig. 4, we necessarily had to reduce the NVAE depth to roughly match its latent dimensionality to other models. Specifically, the NVAE in Fig. 4 had  $2/3/6$  latent groups at top/mid/bottom with  $d = 1$ . This configuration resulted in latent dimensionality of  $D = (2 \cdot 2^2 + 3 \cdot 4^2 + 6 \cdot 8^2) \times 1 = 440$ , which is roughly comparable to other models at  $D = 420$ . See Table 4 for a comparison between cNVAE and the alternative models used in this study.
9.4 Model training details
We generated 750,000 samples for each ROFL category, with a 600,000/75,000/75,000 split for train/validation/test datasets. Model performance (reconstruction loss, NELBO) was evaluated using the validation set (Fig. 14 and 15). Similarly, we used the validation set to estimate mutual information between latents,  $z$ , and ground truth factors,  $g$ , as shown in Fig. 3. To evaluate the latent codes, we used the validation set to train linear regressors to predict  $g$  from  $z$ . We then test the performance using the test set. This includes results presented in Fig. 4 and Fig. 5.
KL annealing and balancing. We annealed the KL term during the first half of the training, which is known to be an effective trick in training VAEs [52, 72, 132, 133]. Additionally, during the annealing period, we employed a KL balancing mechanism which ensures an equal amount of information is encoded in each hierarchical latent group [52, 134, 135].
Gradient norm clipping. During training, we occasionally encountered updates with large gradient norms. Similar to vdvae [71], we used gradient clipping, with an empirically determined clip value of 250. However, we did not skip the updates with large gradients, as we found that the model usually gets stuck and does not recover after skipping an update.
Training hyperparameters. We used batch size of 600, learning rate of 0.002, and trained models for 160 epochs (equivalent to  $160k$  steps). Each training session took roughly 24 hours to conclude on a single Quadro RTX 5000 GPU. Similar to NVAE, we used the AdaMax optimizer [136] with a cosine learning rate schedule [137] but without warm restarts.
9.5 Choosing  $\beta$  values for different figures
It is known that different choices of the  $\beta$  value in the VAE loss (Table 2) will result in different model properties [85]. Therefore, while we scanned across many  $\beta$  values spanning two orders of magnitude, we displayed results for certain betas to best demonstrate our points.
To ensure a fair model comparison in the results shown in Figs. 3, 4, 11 and 12, we selected  $\beta$  values that maximized the overall untangling score for each architecture (cNVAE,  $\beta = 0.15$ ; VAE,  $\beta = 1.5$ ). This can be seen in Fig. 5, which presents DCI scores over the entire  $\beta$  spectrum. The top row of Fig. 5 displays informativeness (= untangling) scores, revealing an inverted U shape, consistent with previous work [85]. Different architectures peak at different  $\beta$  values.
For the neuron shown in Fig. 8b, we found that different  $\beta$  values lead to qualitatively similar outcomes, that is, extreme sparsity in cNVAE feature importances, in sharp contrast to that of VAE. We deliberately chose to show a larger  $\beta = 5$  for VAE because previous work by Higgins et al. [16] suggested that increasing  $\beta$  values increases neuronal alignment, which we did not observe here. In contrast, even a very small  $\beta = 0.01$  for the cNVAE results in a large alignment. This result (paired with other observations) suggests that brain alignment, as we measured it here, emerges due to the architecture rather than from large  $\beta$  values alone—although we do observe some dependence on  $\beta$ , so ultimately a combination of both architecture and loss is important (but mostly architecture).
9.6 Details on latent code evaluation
Each ROFL sample  $\pmb{x} \in \mathbb{R}^{2 \times N \times N}$  is rendered from a low dimensional ground truth vector  $\pmb{g} \in \mathbb{R}^K$ , where  $N$  is the number of pixels (stimulus size) and  $K$  is the true underlying dimensionality of the data (i.e., number of independent ground truth factors). In other words, the data lies on a curved  $K$ -dimensional manifold embedded in a  $2N^2$ -dimensional space.
In the present work, we worked with relatively small stimuli with  $N = 17$ . This choice allowed us to train more models and explore a larger combination of hyperparameters. We chose an odd value for  $N$  so that there was a pixel at the exact center, which served as the center of gaze or the fixation point (Fig. 1a). The true dimensionality (number of generative factors) of different ROFL categories is given in Table 1. For example, it is  $K = 11$  for fixate-1.
We trained unsupervised models by feeding them raw samples  $\pmb{x}$  and asking them to produce a reconstruction  $\hat{\pmb{x}}$ . The models accomplish this by mapping their inputs to latent codes  $z$  during inference (red pathways in Fig. 2), and mapping  $z$  to  $\hat{\pmb{x}}$  during generation (blue pathways in Fig. 2). Note that for all models we have  $z \in \mathbb{R}^{420}$  (see Table 4).
In sum, our approach consists of four steps:
1. Data generation:  $g \to \boxed{\mathrm{ROFL}} \to x$
2. Unsupervised training:  $\pmb{x} \rightarrow \boxed{\text{Model}} \rightarrow \hat{\pmb{x}}$
3. Inference:  $x \to \boxed{\text{Model}} \to z$
4. Evaluation:  $z \leftrightarrow g$ ?
Within this framework, we investigated the relationships between data generative factors  $\pmb{g}$  and model-learned latent codes  $z$  by using the DCI metrics, which we discuss in detail below.
9.7 Metrics & Baselines
We trained a variety of unsupervised models (Table 2) and evaluated them based on (1) their ability to represent the generative variables of the ROFL stimuli; and (2) the neural alignment of their representation. This resulted in a total of five different metrics used to evaluate model performance. Below, we describe the different metrics used for each.
1: untangling & disentangling of data generative factors. To evaluate the latent codes, we employed the DCI framework [22] and made use of the fact that ROFL provides ground truth factors for each data sample. DCI is comprised of three different metrics, Disentanglement, Completeness, and Informativeness, which we will discuss in detail below.
Figure 9: Suppose we train a regression model to predict ground truth factors  $g_{j}$  (blue) from latents  $z_{i}$  (grey). Entanglement (left) refers to a situation in which one latent variable contains information about many generative factors. An overcomplete code (middle) is one in which a single generative factor is predicted from many latent variables. A one-to-one relationship (right) is the ideal outcome.
#2: explaining variability in biological neurons & brain alignment. We used ridge regression to predict macaque MT neuron responses from model latents. To measure the alignment between different latent codes and biological neurons, we used two standard metrics: linear predictive power and brain-alignment score defined in Fig. 8a.
9.7.1 The DCI framework
In this section, we provide more details about the DCI framework [22] and discuss how it was used to evaluate the disentangling performance of our models. The core motivation behind DCI is that a good latent code ideally disentangles as many factors of variation as possible while discarding as little information as possible. In other words, we consider a latent code to be "good" as long as:
1. It can predict the variance of data generative factors, and
2. It exhibits a one-to-one relationship with the generative factors (Fig. 9).
To make this idea concrete, the starting point is to train a regressor  $f_{j}$  (e.g., lasso) that predicts a given data generative factor  $g_{j}$  from a latent code  $z$ . The performance of  $f$  in predicting  $g$  (measured using e.g.,  $R^2$ ) is referred to as the Informativeness score. The notion of Informativeness, if measured using a linear regressor, is very similar to Untangling [20].
As the next step, the weights of  $\pmb{f}$  are used to construct a matrix of relative importances  $R$ , where  $R_{ij}$  denotes the relative importance of  $z_i$  in predicting  $g_j$ . For example,  $R$  can be the magnitude of lasso coefficients. The matrix of relative importance is then used to define two pseudo-probability distributions:
$$
P _ {i j} = \frac {R _ {i j}}{\sum_ {k} R _ {i k}} = \text {p r o b a b i l i t y} z _ {i} \text {b e i n g} g _ {j} \tag {15}
$$
$$
\tilde {P} _ {i j} = \frac {R _ {i j}}{\sum_ {k} R _ {k j}} = \text {p r o b a b i l i t y} g _ {j} \text {b e i n g p r e d i c t e d b y} z _ {i} \tag {16}
$$
These distributions are then used to define the following metrics, which complete the DCI trinity:
$$
D i s e n t a n g l e m e n t: \quad D _ {i} = 1 - H _ {K} \left(P _ {i .}\right); \quad \text {w h e r e} \quad H _ {K} \left(P _ {i .}\right) = - \sum_ {j} P _ {i j} \log_ {K} P _ {i j} \tag {17}
$$
$$
\text {C o m p l e t e n e s s}: \quad C _ {j} = 1 - H _ {D} (\tilde {P} _ {j}); \quad \text {w h e r e} \quad H _ {D} (\tilde {P} _ {j}) = - \sum_ {i} \tilde {P} _ {i j} \log_ {D} \tilde {P} _ {i j} \tag {18}
$$
Here,  $K$  is the number of data generative factors (e.g.,  $K = 11$  for fixate-1, Table 1), and  $D$  is the number of latent variables ( $D = 420$  for the models considered in our paper, Table 4).
Intuitively, if latent variable  $z_{i}$  contributes to predicting a single ground truth factor, then  $D_{i} = 1$ . Conversely,  $D_{i} = 0$  when  $z_{i}$  equally contributes to the prediction of all ground truth factors. If the ground truth factor  $g_{j}$  is being predicted from a single latent variable, then  $C_{j} = 1$  (a complete code).
On the other hand, if all latents contribute to predicting  $g_{j}$ , then  $C_{j} = 0$  (maximally overcomplete). Finally, the overall disentanglement or completeness scores are obtained by averaging across all latents or generative factors.
In the DCI paper, Eastwood and Williams [22] used lasso or random forest as their choice of regressors  $f$ . Initially, we experimented with lasso but found that the lasso coefficient had a strong impact on the resulting DCI scores. To mitigate this, we used linear regression instead and estimated the matrix of relative importance using scikit-learns permutation importance score (sklearn.inspection.permutation Importance), which measures how much performance drops based on the shuffling of a given feature.
9.7.2 Brain alignment score
As depicted in Fig. 8a, the alignment score emphasizes the degree of correspondence between latent variables and neurons. A sparse relationship between latents and neurons suggests a higher alignment, indicating that the model's representation closely mirrors the neuron's representational "form" [16]. This notion of alignment is intimately tied to the "completeness" score discussed above. See Fig. 10.
Figure 10: Consider training a regressor to predict MT neuron responses (green) from latent variables  $z_{i}$  (grey). Beyond accurate predictions, we assess the alignment between  $z$  and MT based on the sparsity of the resulting latent-to-neuron relationships. One possibility is that many latents contribute to predicting the neuron in a dense fashion. This would suggest that the neuron's function is diffusely represented within the latents, hinting that the representational form of the latents is ultimately different from that of MT, perhaps through rotation or affine transformations. Conversely, when few latents effectively predict MT, it suggests a similar representational form, laying the foundation for our rationale in defining our alignment score (Fig. 8a).
9.7.3 Baselines
In Fig. 4 we compared the cNVAE with alternative models introduced in Table 2. We also included an NVAE with matched number of latents, but necessarily fewer hierarchical levels (see section 9.3 for more details). To demonstrate the difficulty of our untangling task, we included two simple baselines: Raw data and a linear code based on the first 420 dimensions of data principal components (PCA). Both Raw and PCA baselines performed poorly, only weakly untangling self-motion velocity ( $\sim 55\%$ ) and object velocity in  $x$  and  $y$  directions ( $47\%$ ). This result demonstrates the importance of nonlinear computation in extracting good representations.
10 Additional Results