SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 4: Hierarchical VAE untangles underlying factors of variation in data. The linear decodability of ground truth factors (x-axis) from different latent codes is shown. Untangling scores averaged across all ground truth factors are cNVAE = 0.898, NVAE = 0.639, VAE = 0.548, cNAE = 0.456, AE = 0.477, PCA = 0.236, and Raw = 0.235. For variational models, the best performing  $\beta$  values were selected: cNVAE,  $\beta = 0.15$ ; VAE,  $\beta = 1.5$  (see Supplementary section 9.5 for more details).
Disentanglement: the cNVAE produces more disentangled representations. The pursuit of disentanglement in neural representations has garnered considerable attention [23, 85, 92-100]. In particular, Locatello et al. [19] established that learning fully disentangled representations is fundamentally impossible without inductive biases. Prior efforts such as  $\beta$ -VAE [85] demonstrated that increasing the weight of the KL loss (indicated by  $\beta$ ) promotes disentanglement in VAEs. More recently, Whittington et al. [92] demonstrated that simple biologically inspired constraints such as non-negativity and energy efficiency encourage disentanglement. Here, we demonstrate that another biological inductive bias, hierarchy in the latent space, will promote disentanglement of the latent representations learned by VAEs.
To evaluate the role of hierarchy, we adopted the DCI framework [22] which offers a well-rounded evaluation of latent representations. The approach involves training a simple decoder (e.g., lasso regression) that predicts data gen
erative factors  $g$  from a latent code  $z$ ; followed by computing a matrix of relative importances (e.g., based on lasso weights) which is then used to evaluate different aspects of the code quality: Informativeness—measures whether  $z$  contains easily accessible information about  $g$  (similar to untangling from above). Disentanglement—measures whether individual latents correspond to individual generative factors. Completeness—measures how many  $z_i$  are required to capture any single  $g_j$ . If a single latent contributes to  $g_j$ 's prediction, the score will be 1 (complete). If all latent variables equally contribute to  $g_j$ 's prediction, the score will be 0 (maximally overcomplete). Note that "completeness" is also referred to as "compactness" [23]. See Fig. 9 and Supplementary
Figure 5: Evaluating the learned latent codes using the DCI framework [22]. Larger values are better for all metrics. Note that informativeness is closely related to untangling [20, 21]. See also Fig. 9.
section 9.7.1 for more details, ref. [101] for a review, and ref. [102] for a recent extension of the DCI framework.
We follow the methods outlined by Eastwood and Williams [22] with two modifications: (1) we replaced lasso with linear regression to avoid the strong dependence on the lasso coefficient that we observed, and (2) we estimate the matrix of relative importances using a feature permutation-based algorithm (sklearn.inspection.permutation Importance), which measures the relative performance drop that results from shuffling a given latent.
We found that cNVAE outperforms competing models across all metrics for a broad range of  $\beta$  values (Fig. 5). The observed pattern of an inverted U shape is consistent with previous work [85], which suggests that there is an optimal  $\beta$  that can be empirically determined. In this case, cNVAE with  $\beta = 0.5$  achieved the best average DCI score. Further, we found that VAEs lacking hierarchical structure learn highly overcomplete codes, such that many latents contribute to predicting a single ground truth factor. In conclusion, the simple inductive bias of hierarchy in the latent space led to a substantial improvement in VAE performance across all components of the DCI metric.
Brain-alignment: the cNVAE aligns more closely with MT neurons. To evaluate the performance of models in predicting neuronal activity in response to motion stimuli, we used an existing dataset of  $N = 141$  MT neurons recorded while presented with random dot kinematograms representing smoothly changing combinations of optic flow velocity fields [103, 104]. A subset of these neurons ( $N = 84$ ) are publicly available on craans.org, and were recently used in Mineault et al. [34] that we compare to.
To measure neuronal alignment, we first determined the mapping from each model's latent representation to MT neuron responses (binned spike counts, Fig. 6a). Here, the latent representation is defined as the mean of predicted Gaussian distributions for VAEs, and the bottleneck activations for AEs. We learn this linear latent-to-neuron mapping using ridge regression. Figure 6b shows the average firing rate of an example neuron along with model predictions. Because sensory neurons have a nonzero response latency, we determined each neuron's optimal response latency, which maximized cross-validated performance. The resulting distribution of best-selected latencies (Fig. 6c) peaked around  $100~ms$ : consistent with known MT latencies [103]. We also empirically optimized ridge coefficients to ensure each neuron has its best fit. Figure 6d shows that the models capture the receptive field properties of MT neurons as measured by the spike-triggered average stimulus. To evaluate performance, we follow methods established by Mineault et al. [34]: whenever repeated trials were available, we report Pearson's  $R$  on that held-out data, normalized by maximum explainable variance [105]. When repeats were not available, we performed 5-fold cross-validation and reported the held-out performance using Pearson's  $R$  between model prediction and spike trains.
Figure 6: (a) Experimental setup form [103, 104]. (b) Both models explain MT neural variability well. (c) Distribution of best estimated latencies. (d) Spike-triggered averages (STA) are shown.
Evaluating brain alignment. We use two measures of brain alignment: the success at predicting the neural response (Pearson's  $R$ , Fig. 7, Table 3); and, the "alignment" between neurons and individual model latents (Fig. 8, [16]). These mirror the untangling and completeness metrics described above (more details are provided below).
Figure 7: All models (pretrained on fixate-1) perform comparably in predicting MT neuron responses. Dashed line corresponds to the previous state-of-the-art on this data [106].
Figure 8: Hierarchical models (cNVAE, cNAE) are more aligned with MT neurons since they enable sparse latent-to-neuron relationships. (a) Alignment score measures the sparsity of permutation feature importances.  $a_{i} = 0$  when all latents are equally important in predicting neuron  $i$ ; and,  $a_{i} = 1$  when a single latent predicts the neuron. (b) Feature importances are plotted for an example neuron (same as in Fig. 6b). cNVAE  $(\beta = 0.01)$  predicts this neuron's response in a much sparser manner compared to non-hierarchical VAE  $(\beta = 5)$ . Supplementary section 9.5 contains a discussion of our rationale in choosing these  $\beta$  values. (c) Alignment across  $\beta$  values, and autoencoders (ae).
All models predict MT neuron responses well. After training a large ensemble of unsupervised models on fixate-1 and learning the neural mapping, we found that both hierarchical (cNVAE & cNAE) and non-hierarchical (VAE & AE) variants had similar ability to predict neural responses (Fig. 7). The performance did depend on the loss function itself, with the variational loss outperforming simple autoencoder reconstruction loss (Table 3).
Hierarchical VAEs are more aligned with MT neurons. We next tested how these factors affect neural alignment, i.e., how closely neurons are related to individual latents in the model. Figure 8a demonstrates what we mean by "alignment": a sparse latent-to-neuron relationship means larger alignment, indicative of a similar representational "form" [16]. See Fig. 10 for an illustration of this idea. To formalize this notion, we use feature permutation importance (described above), applied to the ridge regression models. This yields a 420-dimensional vector per neuron. Each dimension of this vector captures the importance of a given latent variable in predicting the responses of the neuron. We normalize these vectors and interpret them as the probability of importance. We then define alignment score  $a_{i}$  of neuron  $i$  as  $a_{i} = 1 + \sum_{k=1}^{K} p_{ik} \log_{K} p_{ik}$ , where  $p_{ik}$  is interpreted as the importance of  $k$ -th latent variable in predicting neuron  $i$  (Fig. 8a). This concept is closely related to the "completeness" score from the DCI framework as discussed above.
Table 3: Both cNVAE and VAE perform well in predicting MT neuron responses, surpassing previous state-of-the-art models by more than a twofold improvement. Moreover, the clear gap between fixate-1 and other categories highlights the importance of pretraining data [107].
Model	Pretraining dataset	Performance, R (μ ± se; N = 141)
β = 0.5	β = 0.8	β = 1	β = 5
cNVAE	fixate-1	.506 ± .018	.517 ± .017	.494 ± .018	.486 ± .016
fixate-0	.428 ± .018	.450 ± .019	.442 ± .019	.469 ± .018
obj-1	.471 ± .018	.465 ± .018	.477 ± .017	.468 ± .018
VAE	fixate-1	.508 ± .019	.481 ± .018	.494 ± .018	.509 ± .018
cNAE	fixate-1	.476 ± .018
AE	fixate-1	.495 ± .019
CPC [108]	AirSim [75]	.250 ± .020 (Mineault et al. [34])
DorsalNet	AirSim [75]	.251 ± .019 (Mineault et al. [34])
For almost all  $\beta$  values, the cNVAE exhibited a greater brain alignment than non-hierarchical VAE (Fig. 8c; cNVAE > VAE, paired  $t$ -test; see Fig. 16 and Table 5). Similarly, for the autoencoders, we found that the hierarchical variant outperformed the non-hierarchical one (cNAE > AE). Based on these observations, we conclude that higher brain alignment is primarily due to hierarchical latent structure. However, note that hierarchy in the traditional sense did not matter: all these models had approximately the same number of convolutional layers and parameters.
Factors leading to brain-alignment. To test the effect of the training dataset (i.e., category of ROFL) on model performance, we trained cNVAE models using fixate-0, fixate-1, and obj-1 categories (Table 1), while also exploring a variety of  $\beta$  values. We found that fixate-1 clearly outperformed the other two ROFL categories (Table 3), suggesting that both global (e.g., self-motion) and local (e.g., object motion) sources of variation are necessary for learning MT-like representations. The effect of loss function was also visible: some  $\beta$  values led to more alignment. But this effect was small compared to the effect of hierarchical architecture (Fig. 8c).
We introduced a new framework for understanding and evaluating the representation of visual motion learned by artificial and biological neural networks. This framework provides a way to manipulate causes in the world and evaluate whether learned representations untangle and disentangle those causes. In particular, our framework makes it possible to test the influence of architecture (Fig. 2), loss function (Table 2), and training set (Table 1) on the learned representations, encompassing 3 out of the 4 core components of a recently proposed neuroconnectionist research programme [41]. Our framework brings hypothesis-testing to understand [biological] neural processing of vision and provides an interpretive framework to understand neurophysiological data.
The goal of the present work was to establish our framework and demonstrate its potential. To this end, we made several simplifying choices, such as training on individual flow frames rather than time-evolving videos. We provide a detailed discussion of study limitations in Supplementary section 8. Future work will address these by rendering images in simulations and using image-computable models, incorporating real eye-tracking and scene data in ROFL [83, 109], testing our approach on more data from other brain areas such as MST [110, 111], and using more sophisticated methods to measure representational alignment between ANNs and brains [112-115].
Conclusion. We used synthetic data to test how causal structure in the world affects the representations learned by autoencoder-based models and evaluated the learned representations based on how they represent ground truth factors and how well they align with biological brains. We found that a single inductive bias, hierarchical latent structure, leads to desirable representations and increased brain alignment.
6 Code & Data
Our code and model checkpoints are available here: https://github.com/hadivafaii/ROFL-cNVAE.
This work was supported by NSF IIS-2113197 (HV and DAB), NSF DGE-1632976 (HV), and NIH R00EY032179 (JLY). We thank our anonymous reviewers for their helpful comments, and the developers of the software packages used in this project, including PyTorch [116], NumPy [117], SciPy [118], scikit-learn [119], pandas [120], matplotlib [121], and seaborn [122].
[1] Hermann Von Helmholtz. Handbuch der physiologischen Optik. Vol. 9. Voss, 1867.
[2] Ibn al-Haytham. Book of optics (Kitab Al-Manazir). 1011-1021 AD.
[3] David Mumford. "On the computational architecture of the neocortex: II The role of cortico-cortical loops". In: Biological Cybernetics 66.3 (1992), pp. 241-251. DOI: 10.1007/BF00198477.
[4] Tai Sing Lee and David Mumford. "Hierarchical Bayesian inference in the visual cortex". In: JOSA A 20.7 (2003), pp. 1434-1448. DOI: 10.1364/JOSAA.20.001434.
[5] Rajesh PN Rao and Dana H Ballard. "Predictive coding in the visual cortex: a functional interpretation of some extra-classical receptive-field effects". In: Nature Neuroscience 2.1 (1999), pp. 79-87. DOI: 10.1038/4580.
[6] David C Knill and Alexandre Pouget. "The Bayesian brain: the role of uncertainty in neural coding and computation". In: Trends in Neurosciences 27.12 (2004), pp. 712-719. DOI: 10.1016/j.tins.2004.10.007.
[7] Alan Yuille and Daniel Kersten. "Vision as Bayesian inference: analysis by synthesis?" In: Trends in Cognitive Sciences 10.7 (2006), pp. 301-308. DOI: 10.1016/j.tics.2006.05.002.
[8] Karl Friston. “A theory of cortical responses”. In: Philosophical transactions of the Royal Society B: Biological Sciences 360.1456 (2005), pp. 815–836. DOI: 10.1098/rstb.2005.1622.
[9] Andy Clark. "Whatever next? Predictive brains, situated agents, and the future of cognitive science". In: Behavioral and brain sciences 36.3 (2013), pp. 181-204. DOI: 10.1017/S0140525X12000477.
[10] Peter Dayan et al. "The Helmholtz machine". In: Neural Computation 7.5 (1995), pp. 889-904. DOI: 10.1162/neco.1995.7.5.889.
[11] Diederik P Kingma and Max Welling. "Auto-encoding variational bayes". In: (2014). arXiv: 1312.6114v11 [stat.ML].
[12] Danilo Jimenez Rezende et al. "Stochastic backpropagation and approximate inference in deep generative models". In: International Conference on Machine Learning. PMLR. 2014, pp. 1278-1286. URL: https://proceedings.mlr.press/v32/rezende14.html.
[13] Lukas Schott et al. "Towards the first adversarially robust neural network model on MNIST". In: International Conference on Learning Representations. 2019. URL: https://openreview.net/forum?id=S1EH0sC9tX.
[14] Ilker Yildirim et al. "Efficient inverse graphics in biological face processing". In: Science Advances 6.10 (2020), eaax5979. DOI: 10.1126/sciadv.aax5979.
[15] Katherine R Storrs et al. "Unsupervised learning predicts human perception and misperception of gloss". In: Nature Human Behaviour 5.10 (2021), pp. 1402-1417. DOI: 10.1038/s41562-021-01097-6.
[16] Irina Higgins et al. "Unsupervised deep learning identifies semantic disentanglement in single inferotemporal face patch neurons". In: Nature Communications 12.1 (2021), p. 6456. DOI: 10.1038/s41467-021-26751-5.
[17] Joseph Marino. "Predictive coding, variational autoencoders, and biological connections". In: Neural Computation 34.1 (2022), pp. 1-44. DOI: 10.1162/neco_a_01458.
Supplementary material for:
8 Study Limitations & Considerations
We established a synthetic data framework that allows hypothesis generation and testing for neural processing of motion. While our paper opens up many interesting venues for future exploration, it is necessarily simplified, as it focuses on establishing our framework and demonstrating its potential. As such, it currently has several limitations:
First, our simulation generates velocity fields, rather than full spatiotemporal movies, which allows our model to avoid mimicking the complexities of motion extracted by the early visual pathway [123, 124]. This strategy also allows for direct comparison with recorded neural data in MT and MST using random dot kinematograms [103, 104, 106]. However, a more complex model would be necessary to explain responses of neurons earlier in the visual pathway, such as V1, which would require a pixel-computable model as in previous work (e.g., DorsalNet [34]). Likewise, our  $>2\times$  improvement in performance in explaining MT neural data over DorsalNet is likely due in part to their network trained to perform this extraction from spatiotemporal movies, which was not necessarily equivalent to the random dot velocity fields used in the experiments. Thus, it is an open question whether a hierarchical VAE trained and tested on video stimuli would better align to neural data than the model from Mineault et al. [34]. Future work in this space will involve rendering images in simulations and using image-computable models for a fair comparison.
Second, we chose relatively simple environments and simple fixation rules to generate our optic flow fields and avoided the true complexity of 3-D natural environments and their interaction with eye movements and self-motion as has recently been measured [83, 109]. The simplification of our simulation still demonstrates the importance of including such elements in understanding neural representations and provides a framework for incorporating real eye-tracking and scene data [83, 109] into future work with ROFL.
Finally, we only tested neural alignment on one experimental paradigm using neurons in area MT, which leaves the question of whether this is a general principle of brain computation. Addressing this requires testing our approach on more data from other brain areas, such as MST. Based on previous work [106, 110], we expect that hierarchical computation is even more necessary for MST, which is an open question to address in future work.
Interpreting brain-alignment. We measured the alignment between ANN models and MT neurons using both linear predictive power (Fig. 7) and an alternative measure of alignment that is sensitive to the sparsity of latent-to-neuron relationships ("alignment-score", Fig. 8a). Linear regression has been most commonly used to measure similarity, or alignment, between pairs of representations [30, 31, 34, 36], but often results in degenerate [30, 36, 38] and unreliable [125] measures of representational alignment. Consistent with this, our application of linear regression found that it was not effective in differentiating between models: although the cNVAE produced the single best model in terms of neuronal prediction, we found that both hierarchical and non-hierarchical VAEs performed similarly in predicting MT neuron responses (Fig. 7).
In contrast, the alignment score (Fig. 8a) was much more consistent in distinguishing between models (see Fig. 16), and revealed that hierarchical models (both cNVAE and cNAE) had significantly sparser latent-to-neuron relationships. The alignment score measures whether a model has learned a similar representational "form" to the brain, which would enable the sparsity of latent-to-neuron relationships. This concept is closely related to the "completeness" score from the DCI framework [22]. The alignment score shown in Fig. 8a was also used in Higgins et al. [16], although they used the magnitude of nonzero coefficients as their feature importance (under lasso regression). However, we note that this alignment score also has limitations, and for example, it is not a proper metric in the mathematical sense [112]. Future work will consider more sophisticated metrics for brain alignment [112-115].
9 Additional Methods
9.1 VAE loss derivation
Suppose some observed data  $\pmb{x}$  are sampled from a generative process as follows:
$$
p (\boldsymbol {x}) = \int p (\boldsymbol {x}, \boldsymbol {z}) d \boldsymbol {z} = \int p (\boldsymbol {x} | \boldsymbol {z}) p (\boldsymbol {z}) d \boldsymbol {z}, \tag {1}
$$
where  $z$  are latent (or unobserved) variables. In this setting, it is interesting to ask which set of latents  $z$  are likely, given an observation  $x$ . In other words, we are interested in posterior inference
$$
p (\boldsymbol {z} | \boldsymbol {x}) \propto p (\boldsymbol {x} | \boldsymbol {z}) p (\boldsymbol {z}). \tag {2}
$$
The goal of VAEs is to approximate the (unknown) true posterior  $p(\boldsymbol{z}|\boldsymbol{x})$  with a distribution  $q(\boldsymbol{z}|\boldsymbol{x};\theta)$ , where  $\theta$  are some free parameters to be learned. This goal is achieved by minimizing the Kullback-Leibler divergence between the approximate posterior  $q$  and the true posterior  $p$ :
Goal: minimize  $\mathcal{D}_{\mathrm{KL}}\Big[q(\pmb {z}|\pmb {x};\theta)\parallel p(\pmb {z}|\pmb {x})\Big].$  (3)
This objective is intractable, but rearranging the terms leads to the following loss that is also the (negative) variational lower bound on  $\log p(\pmb{x})$ :
$$
\mathcal {L} _ {\mathrm {V A E}} = - \mathbb {E} _ {q} \left[ \log p (\boldsymbol {x} | \boldsymbol {z}; \boldsymbol {\theta} _ {d e c}) \right] + \mathcal {D} _ {\mathrm {K L}} \left[ q (\boldsymbol {z} | \boldsymbol {x}; \boldsymbol {\theta} _ {e n c}) \| p (\boldsymbol {z} | \boldsymbol {\theta} _ {d e c}) \right], \tag {4}
$$
where  $\theta_{enc}$  and  $\theta_{dec}$  are the parameters for the encoder and decoder neural networks (see Fig. 2). The first term in Equation 4 is the reconstruction loss, which we chose to be the Euclidean norm of the differences between predicted and reconstructed velocity vectors (i.e., the endpoint error [84]). We will now focus on the second term in Equation 4, the KL term.
9.1.1 Prior, approximate posterior, and the KL term in vanilla VAE
In standard non-hierarchical VAEs, the prior is not parameterized. Instead, it is chosen to be a simple distribution such as a Gaussian with zero mean and unit covariance:
$$
p (\boldsymbol {z} \mid \theta_ {\text {d e c}}) \rightarrow p (\boldsymbol {z}) = \mathcal {N} (\boldsymbol {0}, \boldsymbol {I}). \tag {5}
$$
The approximate posterior is also a Gaussian with mean  $\pmb{\mu}(\pmb{x};\pmb{\theta}_{enc})$  and variance  $\sigma^2 (\pmb {x};\pmb{\theta}_{enc})$  ..
$$
q (\boldsymbol {z} \mid \boldsymbol {x}; \boldsymbol {\theta} _ {e n c}) = \mathcal {N} \left(\boldsymbol {\mu} (\boldsymbol {x}; \boldsymbol {\theta} _ {e n c}), \sigma^ {2} (\boldsymbol {x}; \boldsymbol {\theta} _ {e n c})\right) \tag {6}
$$
As a result, we see that the KL term for a vanilla VAE only depends on encoder parameters  $\theta_{enc}$ .
9.1.2 Prior, approximate posterior, and the KL term in the cNVAE
Similar to the NVAE [52], the cNVAE latent space is organized hierarchically such that latent groups are sampled sequentially, starting from "top" latents all the way down to the "bottom" ones (i.e.,  $z_{1}$  to  $z_{3}$  in Fig. 2). In addition to its hierarchical structure, another important difference between cNVAE and vanilla VAE is that the cNVAE prior is learned from data. Note the "mid" and "bottom" latent groups in Fig. 2, indicated as  $z_{2}$  and  $z_{3}$  respectively: the cNVAE is designed in a way that changing the parameters along the decoder pathway will impact the prior distributions on  $z_{2}$  and  $z_{3}$  (but not  $z_{1}$ ). Note the "h" in Fig. 2, which is also a learnable parameter. In summary, the "top" cNVAE latents (e.g.,  $z_{1}$  in Fig. 2) have a fixed prior distribution similar to vanilla VAEs; whereas, the prior distribution for every other cNVAE latent group is parametrized and learned from data.
More formally, the cNVAE latents are partitioned into disjoint groups,  $z = \{z_{1},z_{2},\ldots ,z_{L}\}$ , where  $L$  is the number of groups. Then, the prior is represented by: