Hierarchical VAEs provide a normative account of motion processing in the primate brain
Authors: Hadi Vafaii
Jacob L. Yates
Daniel A. Butts
Venue: NeurIPS 2023 poster (2023)
Keywords: NeuroAI
VAE
Dorsal stream
Hierarchical Bayesian Inference
TL;DR: We used hierarchical VAEs and a novel simulation to test the hypothesis that brains perform hierarchical Bayesian inference about the state of the world.

Abstract: The relationship between perception and inference, as postulated by Helmholtz in the 19th century, is paralleled in modern machine learning by generative models like Variational Autoencoders (VAEs) and their hierarchical variants. Here, we evaluate the role of hierarchical inference and its alignment with brain function in the domain of motion perception. We first introduce a novel synthetic data framework, Retinal Optic Flow Learning (ROFL), which enables control over motion statistics and their causes. We then present a new hierarchical VAE and test it against alternative models on two downstream tasks: (i) predicting ground truth causes of retinal optic flow (e.g., self-motion); and (ii) predicting the responses of neurons in the motion processing pathway of primates. We manipulate the model architectures (hierarchical versus non-hierarchical), loss functions, and the causal structure of the motion stimuli. We find that hierarchical latent structure in the model leads to several improvements. First, it improves the linear decodability of ground truth variables and does so in a sparse and disentangled manner. Second, our hierarchical VAE outperforms previous state-of-the-art models in predicting neuronal responses and exhibits sparse latent-to-neuron relationships. These results depend on the causal structure of the world, indicating that alignment between brains and artificial neural networks depends not only on architecture but also on matching ecologically relevant stimulus statistics. Taken together, our results suggest that hierarchical Bayesian inference underlines the brain's understanding of the world, and hierarchical VAEs can effectively model this understanding.

Hierarchical VAEs provide a normative account of motion processing in the primate brain
Hadi Vafaii vafail iffaiion 1ummd edu
Jacob L. Yates? yates@berkeley, edu
Daniel Butts dab@umd edu
University of Maryland, College Park
UC Berkeley
Abstract
The relationship between perception and inference, as postulated Helmholtz in the 19th century, is paralleled in modern machine learning by generative models like Variational Autoencoders (VAEs) and their hierarchical variants. Here, we evaluate the role of hierarchical inference and its alignment with brain function in the domain of motion perception. We first introduce novel synthetic data framework, Retinal Optic Flow Learning (ROFL), which enables control over motion statistics and their causes. We then present a new hierarchical VAE and test it against alternative models on two downstream tasks: (i) predicting ground truth causes of retinal optic flow (e.g., self-motion); and (ii) predicting the responses of neurons in the motion processing pathway of primates. We manipulate the model architectures (hierarchical versus non-hierarchical), loss functions, and the causal structure of the motion stimuli. We find that hierarchical latent structure in the model leads to several improvements. First, it improves the linear decodability of ground truth factors and does so in sparse and disentangled manner. Second, our hierarchical VAE outperforms previous state-of-the art models in predicting neuronal responses and exhibits sparse latent-to-neut depend on the causal structure of the world, indicating that alignment between brains and artificial neural networks depends not only on architecture but also on matching ecologically relevant stimulus statistics. Taken together, our results suggest that hierarchical Bayesian inference underlines the brain's 's understanding of the world, and hierarchical VABs can effectively model this understanding.
relationships. These results
Introduction
Intelligent interactions with the world require representation of its underlying composition. This inferential process has long been postulated to underlie human perception (1.9). and paralleled in modern machine lcarning by generative models 17). which learn latent representations of their sensory inputs. The question of what constitutes "good" representation has no clear answer but several desirable features have been proposed. In the field of neuroscience, studies tocused on object recognition have suggested that effective representations "untangle" the various factors of variation in the input, rendering them linearly decodable (20) 21). This intuitive notion of linear decodability has emerged in the machine learning community under different names such as "informativeness" [22) or "explicitness" [23) Additionally, it has been suggested that "disentangled" representations are desirable, wherein distinct, informative factors of variations in the data are separated 24].29] Artificial neural networks (ANNs) are also increasingly evaluated based on their alignment with biological neural processing (8), because the shared goals of ANNs and the
18 19
brain's 's sensory processing [25] 39 40) Such alignment also provides the possibility of gaining
insights into the brain by understanding the operations within an ANN 47
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

In this work, we investigate how the combination of (i) model architecture, (ii) loss function, and (iii) training dataset, affects learned representations. and whether this is related to the brain- -alignmmm of the ANN [4] 44
We focus specifically on understanding the representation of motion because large sections of the visual cortex are devoted to processing motion (34), and the causes of retinal motion (moving objects and self-motion
(48) can be manipulated systematically. Crucially, motion in an image can be described irrespective of the identity and specific visual features that are moving, just as the identity of objects is invariant to how they are moving. This separation of motion and object processing mirrors the division of primate visual processing into dorsal (motion) and ventral (object)
streams 49 5T].
We designed a naturalistic motion simulation based on distributions of ground truth factors cor- responding to the location and depth of objects, motion of these objects, motion of the observer, and observer's direction of gaze (i.e., the fixation point; Fig. [A). We then trained and evaluated an ensemble of autoencoder-basee models using our simulated retinal flow data. We based our evaluation on (1) whether the models untangle and disentangle the ground truth factors in our simulation; and (2) the degree to which their latent spaces could be directly related to neural data recorded in the dorsal stream of primates (area MT).
We introduce a new hierarchical variational autoencoder, the "compressed" Nouveau VAE (cNVAB) (52). The cNVAB exhibited superior performance compared to other models across our multiple evaluation metrics, First, it discovered latent factors that accurately captured the ground truth factors in the simulation in a more disentangled manner than other models. Second, it achieved significant improvements in predicting neural responses compared to the previous state-of-the-arr model (34). doubling the performance, with sparse mapping from its latent space to neural responses.
Taken together, these observations demonstrate the power of the synthetic data framework and show that single inductive bier-hiiralc latent structure-leads many desirable features of representations, including brain alignment.
2 Background & Related Work
Neuroscience and VAES. It has long been argued that perception reflects unconscious inference of the structure of the world constructed from sensory inputs. The concept of "perception as unconscious
inference" has existed since at least the 19th century 21. and more recently inspired Mumford (3] to conjecture that brains engage in hierarchical Bayesian inference to comprehend the world [3] 4
These ideas led to the development of Predictive Coding 9 58), Bayesian Brain Hypothesis 6 (9)655), and Analysis-by- saalyss-bb-Syythesiis (7). collectively suggesting that brains contain an internal
generative model of the world (7 8 66 67). A similar idea underlies modern generative models
17 68 (70,, especially hierarchical variants of VAFs [52] 73..
The Nouveau VAE (NVAE) (52) and very deep VAE (vdvae) (71] demonstrated that deep hier- archical VAEs can generate realistic high-resolution images, overcoming the limitations of their non-hierarchical predecessors. However, neither work evaluated how the hierarchical latent structure changed the quality of learned representations. Additionally, both NVAE and vdvae have an undesir- able property: their convolutional latents result in a latent space that is several orders of magnitude larger than the input space, defeating a main purpose of autoencoders: compression. Indeed, Hazami et al. (74) showed that a tiny subset (around 3%) of the vdvae latent space is sufficient for comparable input reconstruction. Here, we demonstrate that is possible to compress hierarchical VAFs and focus on investigating their latent representations with applications to neuroscience data.
Evaluating ANNs on predicting biological neurons. Several studies have focused on evaluating ANNs on their performance in predicting brain responses, but almost entirely on describing static ("ventral stream")) image processing (30 33 36). In contrast, motion processing (corresponding to the dorsal stream) has only been considered thus far in Mineault et al. [34], who used 3D ResNet ("DorsalNet") to extract ground truth factors about self-motion from drone footage ("AirSim",, (75) in supervised manner. DorsalNet learned representations with receptive fields that matched known features of the primate dorsal stream and achieved state-of-the-art on predicting neural responses on the dataset that we consider here. In addition to our model architecture and training set, a fundamental difference between our approach and Mineault et al. (34) is that they trained their models using direct supervision. As such, their models have access to the ground truth factors at all times.

---

Table I: ROFL categories used in this paper. ground truth factors include fixation point (+2)); velocity of the observer when self-motion is present (+3); and, object position & velocity (+6). Figure P showcases few example frames for each category. The stimuli can rendered at any given spatial scale N, yielding an input shape of 2 x N x N. Here we work with N = 17.
Category
Description
Dimensionality
A moving observer maintains fixation on background point.
fixate-1
11 = 2 3+
In addition, the scene contains one independently moving object.
fixate-0
Same as f fixate-1 but without the object.
5 = 2 3
obj-1
A single moving object, stationary observer.
8 = 2
Here, we demonstrate that it is possible to obtain ground truth factors "for free", in a completely unsupervised manner, while achieving better performance in predicting biological neuronal responses.
core component of a reductionist approach to studying the brain is to characterize neurons based on their selectivity to particular subset of pre-identified
Using synthetic data to train ANNs.
visual "features", usually by presenting sets of "feature-isolating" stimuli (76). In the extreme, stimuli are designed that remove all other features except the one under investigation (77). While these approaches can inform how pre-selected feature sets are represented by neural networks, is often difficult to generalize this understanding more natural stimuli, which are not necessarily well-described by any one feature set. As a result, here we generate synthetic data representing a naturalistic distribution of natural motion stimuli. Such synthetic datasets allow us manipulate the causal structure of the world, in order to make hypotheses about what aspects of the world matter for the representations learned by brains and ANNs (78). Like previous work synthesized textures 15], here we specifically manipulate the data generative structure to contain factors of variation due to known ground truth factors.
3 Approach: Data & Models
Retinal Optic Flow Learning (ROFL.. Our synthetic dataset framework, ROFL,, generates the resulting optic flow from different world structures, self-motion trajectories, and object motion (Fig. D see also (79). )
ROFL can be used to generate naturalistic flow fields that share key elements with those experienced in navigation through 3-D environments. Specifically, each frame contains global patterns that are due
to self-motion, including rotation that can arise due eye or head movement [80 81]. In addition,
local motion patterns can be present due to objects that move independently of the observer [48] The
overall flow pattern is also affected by the observer's direction of gaze (fixation point 82 Fig. TT).
ROFL generates flow vectors that are instantaneous time, representing the velocity across the visual field resulting from the spatial configuration of the scene and motion vectors of self and object. given scene (which can arguably be considered separably (83]
Ignoring the time-evolution dramatically reduces the input space from [3 x H x W T" to [2 x H x W), and allows broader sampling of configurations without introducing changes in luminance and texture. As a result, we explore the role of different causal structures in representation learning in ANNs.
The retinal flow patterns generated by a moving object depend on both the observer's 's self-motion and the rotation of their eyes as they maintain fixation in the world, in addition to the motion of the object itself. For example, Fig. [Tc demonstrates a situation where the observer moving forward, and the object is moving to the right, with different object positions: an object on the left side will have its flow patterns distorted, while an object on the right will have its flow patterns largely unaffected because its flow vectors are parallel with that of the self-motion. In summary, ROFL allows to simulate retinal optic flow with a known ground truth structure driven by object and self-motion.
The compressed NVAE (cNVAE). The latent space of the NVAE is partitioned into groups, where is the number of groups. The latent groups are serially dependent, meaning that the distribution of a given latent group depends the value of the preceding latents,
(21, ..,..., }
3

---

a
b
Object
X
Observer
fixate-O
Z
(0.0,
y
mackground
Solid
obj-1
C
F, 0,as --
F, K,y
F 0,N --2
- 0, -0
= 0,
F K.y -4
F, 0,
Fu== 4, Xaw
Fa= 2, Kay
1= Xay =4
F== Xay
f==1, Xay 44
F.= Xa=
f,=4, X
Figure 1: Retinal Optic Flow Learning (ROFL): simulation platform for synthesizing naturalistic optic flow patterns. (a) The general setup includes moving or stationary observer and a solid background, with optional moving object(s) in the scene. More details provided in the appendix (section [3). (b) Example frames showcasing different categories (see Table for definitions). (c, d) Demonstrating the causal effects of varying single ground truth factor while keeping all others fixed: (c) Xogj. the x component of object position (measured in retinal coordinates, orange), and (d) Fz, the X component of the fixation point (measured in fixed coordinates, gray).
such that the prior is given by p(z) = p(z1) H p(zz kzt), and approximate posterior is given by ((2)0)) (z) I (= æ) (more details in section 9.1 Additionally, different latent groups in the NVAE operate at difffernnt spatial scales (Fig. 2 left), with multiple groups scale. Crucially, such scale-dependent grouping is absent from hierarchicl VAEs (Fig. 2 right).
The eNVAE closely follows the NVAE (52), with one important difference: the original NVAE latent space is convolutional, and ours is not. We modified the sampler layers (grey trapezoids, Fig. 2 such that their receptive field sizes match the spatial scale they operate on. Thus, sampler layers integrate spatial information before sampling from the approximate posterior. The spatial patterns of each latent dimension are then determined by expand modules (yellow trapezoids, Fig. 2). based on deconvolution step. Further details about the processing of the sampler and expand layers provided in Supplementary section 9.2
Our modification of the NVAE serves two purposes. First, it decouples spatial information from the functionality of latent variables, allowing them to capture abstract features that are invariant to particular spatial locations. Second, it has the effect of compressing the input space into lower-dimensional latent code. We explain this in more detail in Supplementary section 9.3
Our model has the following structure: 3 latent groups operating at the scale of 2 x 2; 6 groups at the scale of x 4; and 12 groups at the scale of x 8 (Table 4 Fig. 2 Therefore, the model has + 6 + 12 21 hierarchical latent groups in total. Each latent group has 20 latent variables, which results an overall latent dimensionality 21 x 20 = 420. See Table and Supplementary section 9.3 for more details.
Alternative models. We evaluated a range of unsupervised models alongside VAF, including standard (non-hicrarchical) VAEs [I] 12 hierarchical autoencoder with identical architecture as

---

Table 2: Model details. Here, hierarchical means that there are parallel pathways for information to flow from the encoder to the decoder (Fig. 2 which is slightly different from the conventional notion. For variational models, this implies hierarchical dependencies between latents in a statistical sense (71). This hierarchical dependence is reflected in the KI. term for the cNVAF, where T. is the number of hierarchical latent groups. See Supplementary section 993 for more details and section 99. for derivation. All models have equal # of latent dimensions (420, see Table 4. approximately the same of convolutional layers, and of parameters (~ 24 M). EPE, endpoint error.
Model
Architecture
Loss
Kullback-Leibler term (KL)
KT. Σ T ( [T.,]] where KLg DKL (= =c() (((e
cNVAE
Hierarchical
EPE +B x KL
VAE
Non-hierarchical
EPE +.3 KL
KL DKL ((2(22) [x(2)
cNAE
Hierarchical
EPE
AE
kon-hierarchical
EPE
the cNVAE but trained only with reconstruction loss (cNAE), and autoencoder (AE) counterpart for the VAE (Table All models had the same latent dimensionality (Table and approximately the same number of parameters and convolutional layers. We used endpoint error as our measure of reconstruction loss, which is the Euclidean norm of the difference between actual and reconstructed flow vectors. This metric works well with optical flow data (84).
Model representations. We define a model's internal representation to be either the mean of each Gaussian for variational models (i.e., samples drawn from q (z [x) at zero temperature), or the bottleneck activations for au- toencoders. For hierarchical models (cNVAE, cNAB), concatenate representations across all levels (Table 4
2 x 2
Training details. Models were trained for 160,000 steps at an input scale of 17x 17, requiring slightly over day Quadro RTX 5000 GPUs. Please refer to Supplementary section 9.4 for additional details.
4 x 4
Disentanglement and 3 V-VAEs. A critical decision when optimizing VAEs involves determining the weight assigned to the KL term in the loss function compared to the reconstruction loss. Prior research has demonstrated that modifying a single parameter, denoted as 3, which scales the KL term, can lead the emergence of disen- tangled representations [85] 86). Most studies employing VAEs for image reconstruction typically optimize the stan- dard evidence lower bound (ELBO) loss, where /3 is fixed value of [TI 52 71). However, should be noted that due to the dependence of the reconstruction loss on the input size, any changes in the dimensionality of the input will inevitably alter the relative contribution of the KL term, and thus the "effective" / (85).
8 X 8
C
C
Figure 2: Architecture comparison. Left, compressed NVAE (cNVAE): right, non- hierarchical VAE. We modified the NVAE sampler layer (grey trapezoids) and introduced a deconvolution expand layer (yellow trapezoids). The encoder (inference) and decoder (generation) pathways are depicted in red and blue, re- spectively. residual block; h, trainable parameter; +, feature combination.
Furthermore, Higgins et al. 166] recently established strong correspondence between the generative factors dis- covered by 3- 3-VAEs and the factors encoded by inferotem- poral (IT) neurons in the primate ventral stream. The alignment between these factors and ITT neurons exhibited a linear relationship with the value of (3. In light of these findings, we explicitly manipulate the parameter 3 within a range spanning from 0.01 to 10 to investigate the extent to which our results depend on its value.
5

---

4 Results
Our approach is based on the premise that the visual world contains a hierarchical structure. We use a simulation containing a hierarchical structure (ROFL, described above) and a hierarchical VAE (the eNVAE, above) to investigate how these choices affect the learned latent representations, While we are using relatively simple simulation generated from a small number of ground truth factors, g, we do not specify how g should be represented in our model include g in the loss. Rather, we allow the model to develop its own latent representation in a purely unsupervised manner. See Supplementary section for more details on our approach.
We first consider hierarchical and non-hierarchical VAES trained on the f fixate-1 condition (see [able through- out this work, fixate-1 is used unless stated otherwise). We extracted latent representations from each model and estimated the mutual information (MI) between the representations and ground truth factors such as self-motion, etc. For f fixate-1, each data sample is uniquely determined using 11 ground truth factors (Table D. and the models have latent dimensionality of 420 (Table Thus, the resulting MI matrix has shape 11 420, where each entry shows how much information is contained in that latent variable about a given ground truth factor.
2
828
100ACAA 0000000
100 120 342 363
380
200 228 240
260
700
300
300
340
300
300
400
420
Figure Mutual information between latent variables (x-axis) and ground truth factors (y-axis) is shown for cNVAE (top) and VAE (bottom). Dashed lines indicate 21 hierarchical latent groups of 20 latents each, com- prising a 420-dimensional latent space. These groups operate three different spatial scales, as indicated. In contrast, the VAE latent space lacks such grouping and operates solely at the spatial scale of 2 x 2 (see Fig. and Table 4 for details on model latent configurations).
Functional specialization emerges in the CNVAE. Figure 3 shows the MI ma- trix for the latent space of cNVAE (top) and VAE (bottom). While both models achieved good reconstruction of valida- tion data (Fig. 14) the MI matrix for cN-
VAE exhibits clusters corresponding to distinct ground truth factors at different levels of the hierarchy. Specifically, object-related factors of variation are largely captured at the top 2 x 2 scale, while information about fixation point can be found across the hierarchy, and self-motion is largely captured by 8 x 8 latent groups. In contrast, non-hierarchical VAE has no such structure, suggesting that the inductive bias of hierarchy enhances the quality of latent spaces, which we quantify next.
Evaluating the latent code. To demonstrate the relationship between ground truth factors and latent representations discovered by the cNVAE visible in Fig. 3 we apply metries referred as "untangling" and "ddsentengling". Additionally, in a separate set of experiments, we also evaluate model representations by relating them to MT neuron responses, which we call "brain-alignmen"" We discuss each of these in detail in the following sections.
Untangling: the cNVAE untangles factors of variation. One desirable feature of latent repre-
sentation is whether it makes information about ground truth factors easily (linearly) decodable (20]
27] This concept has been introduced in the context of core object recognition as "untangling"
Information about object identity that is "tangled" in the retinal input is untangled through succes-
sive nonlinear transforms, thus making it linearly available for higher brain regions to extract (20).
This concept is closely related to the "informativeness" metric of Eastwood and Williams 22 and
"explicimess"" metric of Ridgeway and Mozer 23
To assess the performance of our models, we evaluated the linear decodability of the ground truth factors, g, from model latent codes, z. Based on the R²2 scores obtained by predicting g from z using linear regression (Fig. 4.. the cNVAE greatly outperforms competing models, faithfully capturing all ground truth factors. contrast, the non-hierarchical VAE fails to capture object-related variables. Notably, the cNVAE can recover the fixation point location (Fx, Fy) in physical space almost

---

perfectly. The fixation location has a highly nontrivial effect on the flow patterns, and varying it causes both global and local changes in the flow patterns (Fig. Tt).
Furthermore, eNVAE is the only model that reliably captures object position and velocity: especially note Vobj.= (last column in Fig. 4. Inferring object motion from complex optic flow patterns involves two key components. First, the model must extract self-motion from flow patterns. Second, the model must understand how self-motion influences flow patterns globally. Only then can the model subtract self-motion from global flow vectors to obtain object motion. In vision science, this is known as the "flow-pane flow-parsing hypothesis" 88 9I]. Such flow-pal flw-parsing is achieved by the eNVAE but none of the other models. See Supplementary section I1 for further discussion of this result and its implications.
1.0
0.6
PCA
0.4-
0.2
0.0
Vasrr,
Vsetty
Vsait,z
Xaij
Yoo!
Zooj
Vanj.x
Vanj.
Voog
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 eN VAVAE 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 values were selected: eNVAE, B 0.15; VAE, B 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
Informativeness
cNWIE
gamered considerable attention [23] 85 92
WAE
0.5
cNAE
100]. In particular, Locatello et al. 19 estab- lished that learning fully disentangled represen- tations is fundamentally impossible without in- ductive biases. Prior efforts such as 3-VAE [85] demonstrated that increasing the weight of the KL loss (indicated by (3) promotes disentangle- ment 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 represen- tations learned by VAEs.
AE
0.0
Disentanglement
0.5
0.0
Completeness
0.5
0.0
0
0
STO
0.0
0.
000
111
00
80
00
05-
000
B
To evaluate the role of hierarchy, we adopted Figure 5: Evaluating the learned latent codes using the DCI framework (22). Larger values better
the DCI framework (22 which offers 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 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 qual- ity: 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-meante single gj. If a single latent contributes to أ 's prediction, the score will be 1 (complete). If all latent variables equally contribute to gj's prediction, the score will be 0 (maximally overcomplete). Note that "completeness" is also referred to "compacthess"
for all metrics. Note that informativeness is closely related untangling 20 21). See also Fig. 9
how many 2; are required to capture any
23 See Fig. 9 and Supplementary
7

---

section 9.7.] for more details, ref. [101] for a review, and ref. [002] 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 avoid the strong dependence the lasso coefficient that observed, and (2) we estimate the matrix of relative importances using a feature permutation- algorithm ieffearn.inspectinn pet.iionn.emmuttiionn.. immpptttmmpppeeh which measures the relative performance drop that results from shufiling a given latent.
We found that cNVAE outperforms competing models across all metrics for broad range of 3 values (Fig. 5). The observed pattern of an inverted U shape consistent with previous work [85]. which suggests that there is optimal / that can be empirically determined. In this case, cNVAE with 3 = 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 substantial improvement in VAE performance across all components of the DCI metric.
Brain-alignment: the EN aligns more closely with MT neurons. To evaluate the performance of models predicting neuronal activity in response 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 [T03 [04). subset of these neurons
(N 84) are publicly available on crens.org and were recently used in Mineault et al. [34] that we compare to.
To measure neuronal aligument, we first determined the mapping from each model's latent repre- sentation MT neuron responses (binned spike counts, Fig. Fig. 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 G shows the average firing rate of an example neuron along with model predictions. Because sensory neurons have nonzero response latency, we determined each neuron's optimal response latency, which maximized cross-validated performance. The resulting distribution of best-selected latencies (Fig 6 peaked around 100 ms: consistent with known MT latencies [03). We also empirically optimized ridge coefficients to ensure each neuron has its best fit. Figure 6 shows that the models capture the receptive field properties of MT neurons as measured by the evaluate performance, we follow methods established by Mineault et al. [34]: whenever repeated trials were available, we report Pearson's 's R on that held-out data, normalized by maximum explainable variance [105]. When repeats were not available, we performed 5-fold cross-vall coss-validation and reported the held-out performance using Pearson's R between model prediction and spike trains.
average stimulus. To
Spike counts
b
C
LIOKaY (nx)
Macaque
electrophysiology
Firing rate spikes/s]
Count g.
cNVAE, 66. % =5
MT
a
STA from
Spikes
20
model pred.
STAfrom
Moving dots motion stimuli
0.0
1.0
2.0
3.0
4.0
5.0
Time [s]
cNVAE
VAE
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: success at predicting the neural response (Pearson's 's R, Fig. 7 Table 3 and, the "alignment" between neurons and individual model latents (Fig. 8 [T6).. These mirror the untangling and completeness metries described above (more details are provided below).

---

1.0
model
Performance
0.8
CNVAE
0.6
VAE
cNAE
0.4
AE
0.2
0.0
ae
0.01
0.1
0.15
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
1.0
1.5
2.0
5.0
10.0
B
Figure 7: All models (pretrained fixate-1) perform comparably in predicting MT neuron responses. Dashed line corresponds to the previous state-of-the-an on this data [066..
b
K Pix log Pit
MT neuron
cNVAE
Q:
Feature importance
0.4
0.87
T=1
Alignment score
P.k
VAE
0.04
a
0.20
0.02
k
0.00
Latent dimensions
100
200
300
400
C
Latent dimensions
model
Aligmment
<NVAE
0.6
VAE
cNAE
0.4
0.2
0.0
88
0.01
0.1
0.15
0.2
0.4
0.5
0.6
0.7
0.8
0.9
1.0
1.5
2.0
10.0
B
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. when a single latent predicts the neuron. (b) Feature importances are plotted for example neuron (same as in Fig. ()). cNVAE (33 0.01) predicts this neuron's 's response in a much sparser manner compared to non-hierarchical VAE (3 5). Supplementary section contains a discussion of our rationale in choosing these 3 values. (c) Alignment across 3 values, and autoencoders (ae).
= 0 when all latents are equally important in predicting neuron i; and,
All models predict MT neuron responses well. After training a large ensemble of unsupervised models fixate-1 and learning the neural mapping, found that both hierarchical (cNVAE & cNAE) and oon-hierarchiaal (VAE & AE) variants had similar ability to predict neural responses (Fig. 7 The performance did depend the loss function itself, with the variational loss outperform- ing simple autoencoder reconstruction loss (Table 3).
Hierarchical VAEs are more aligned with MT neurons. We next tested how these factors affect neural alignment, i... how closely neurons are related to individual latents in the model. Figure B demonstrates what mean by "alignment": sparse latent-to neuron relationship means larger alignment, indicative similar representational "form" [16]. See Fig. 10 for an illustration this feature permutation importance (described above), applied
idea. To formalize this notion, to the ridge regression models. This yields a 420-dimeensionll vector per neuron. Fach dimension of this vector captures the importance of 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; of neuron a; = 1 Σ Pik log Pik. where Pik is interpreted the importance of k-th latent variable predicting neuron (Fig. (h). 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 thte-off-hheart models by more than twofold improvement. Moreover, the clear gap between f fixate-1 and other categories highlights the importance of pretraining data [107].
Pretraining dataset
Performance, R (n 8E; 141)
Model
= 0.5
= 0.8
8 = 1
8 = 5
fixate-1
.006 + 0018
.17 .017
.494 + .018
.486 + .016
cNVAE
fixate-0
.428 .018
.450 + .019
.442 + .019
.469 + .018
obj-1
.471 + .018
.465 .018
.47 + .017
.468 + .018
VAE
fixate-1
.508 = .019
.481 = .018
494 .018
.509 + .018
cNAE
fixate-1
.476 .018
AE
fixate-1
.455 .019
CPC (108)
AirSim [75]
.50 + .020 (Mineault et al.
34)
DorsalNet
AirSim [75]
.251 + .019 (Mineault et al. 344
For almost all 3 values, the eNVAF exhibited a greater brain alignment than non-hierarchical VAE (Fig. cNVAE VAE, paired t-test; Fig. and T) Similarly, for the autoencoders, found that the hicrarchical variant outperformed the non-hice non-hicrarchical 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-aligument. To test the effect of the training dataset (i.c., category of ROFL) on model performance, we trained cNVAE models using f fixate-0, fixate-1, and obj categories (Table D. while also exploring a variety of 3 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 3 values led to more alignment. But this effect was small compared to the effect of hierarchical architecture (Fig. (c).
5 Discussion
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 [ on the learned representations, encompassing out of the 4 core components of 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 neurophysiologica data.
The goal of the present work was 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 detailed discussion of study limitations in Supplementary section g Future work will address these by rendering images in simulations and using congeccomputalle models, incorporating real eye-tracking and scene data ROFL [83 [09], 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 (152-
Conclusion. We used synthetic data to test how causal structure in the world affects the representa- tions learned by autoencoder-basee 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 single inductive bias, hierarchical latent structure, leads to desirable representations and increased brain alignment.
10

---

Code & Data
Our code and model checkpoints are available here: http://ghhhbbccmmmaaaaaa////ffff
7
mcknowledgments
This work was supported by NSF IIS-2113197 (HV and DAB), NSF DGEE-1329766 (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 Py Torch [116), NumPy [117]
SciPy [18] scikit-learn [19), pandas 120]. matplotlib [121], and seaborn 122
References
[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 20.7 (2003), pp. 1434-1448. DOI: /0141/0.364///// JOSAA. 20. 001434
[5] Rajesh PN Rao and Dana Ballard. "Predictive coding in the visual cortex: a functional interpretation of some extra-classical receptive-fiele field effects". In: Nature Neuroscience 2.1
(1999), pp. 79-87. DOT: 10. 1008/4534550 /488
[6] David C Knill and Alexandre Pouget. "The Bayesian brain: the role of uncertainty in neural coding and computation". In: Trends Neurosciences 27.12 (2004), pp. 712-719. DOr:
10. 1016/j.timm 2004 1004.00.07
[7]
Alan Yuille and Daniel Kersten. "Vision Bayesian inference: analysis by synthesis?" In: Trends in Cognitive Sciences 10.7 (2006), pp. 301-308. DOI: 10. 10t.1006....i. tics 2006 05
002
[8]
Karl Friston. "A theory cortical responses". In: Philosophical transactions of the Royal Society B: Biological Sciences 360.1456 (2005), pp. 815-836. DOI: 10. www.//9tttb2000
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 181-204. DO1: 10. 1017
80140525X1200044
[10] Peter Dayan et al. "The Helmholtz machine", In: Neural Computation 7.5 (1995), 889-
904. DOL: 10. wwe11166cccc... 1.95.7.5.88..
[11] Diederik Kingma and Max Welling. "Auto-encoding variational bayes". In: (2014). arXiv:
1312. 6112.6114v1 [stat. ]
[12] Danilo Jimenez Rezende al. "Stochastic backpropagation and approximate inference in deep generative models". In: International Conference on Machine Learning. PMLR. 2014, pp. 1286. URL: lattps://proceedinggss mlr wtw.:/www..p....eeeeememmmmmeeee watpaa//vvvreeeeeeeeeee html
[13] Lukas Schott MNIST" In: International Conference on Learning Representations. 2019. URL: https
"Towards the first adversarially robust neural network model on
/openreview net/ wwt.:eettooonneee SEHomessssS1S1HH1
[14] llker Yildirim et al. "Efficient inverse graphics in biological face processing". In: Science Advances 6.10 (2020), c9ax5979, Do1: 10. /126/sciade aax5979
[15]
Katherine R Storrs al. "Unsupervised learning predicts human perception and misperception
of gloss". In: Nature Human Behaviour 5.10 (2021), pp. 1402-1417. DOI: 10. 0038//41556
021 01097-6
[16]
Irina Higgins et "Unsupervised deep learning identifies semantic disentanglement in single inferotemporal face patch neurons". In: Nature Communications 12.1 (2021), 6456. DOI:
10. -22------2212-
[17] Joseph Marino. "Predictive coding, variational autoencoders, and biological connections". In:
Neural Computation 34.1 (2022), pp. 1-44, DOT: 10. 1 1162/neco a_ 01458

---

[18]
Irina Higgins et al. "Towards a definition of disentangled representations". In: (2018). arXiv: 1812. 02230 [c]
[19]
Francesco Locatello et al. "Challenging common assumptions in the unsupervised learning of disentangled representations" In: international conference machine learning. PMLR. 2019, pp. 4124 URL: https: http:::://oocceiingss... mlr press/ (atoss/vv///octeell/// locatello19a html
[20] James J DiCarlo and David D Cox. "Untangling invariant object recognition", In: Trends in Cognitive Sciences 11.8 (2007), pp. 333-341. DO1: 10. 1016/j wticsiiiii.i..o...o.ooooooiooooa tics 2007 060
[21]
James J DiCarlo al. "How does the brain solve visual object recognition?" In: Neuron 73.3 (2012), pp. 415-434, DO1: 10. 1016/j neuron. 20100000.10000
[22] Cian Eastwood and Christopher K. I. Williams. "A framework for the quantitative evaluation of disentangled representations"" In: International Conference on Learning Representations.
2018. URL: /Rtp.ttttts:///peeerrriiiiii ottp://///ppeeeeeeeeeeeeeeerrrra -d-
[23] Karl Ridgeway and Michael C Mozer. "Learning Deep Disentangled Embeddings With the F-Statistic Loss". " In: Advances in Newral Information Processing Systems. Vol. 31, Curran Associates, Inc., 2.18. URL: https: //tp:::///pppaeeee nips wwwpappp.p.p.eeeepcpeep filfffiess wttpsffllsssppee///eeeeea paper httl://8244500228666634466666613
Abstraatttt..................aat html
[24] Yoshua Bengio et al. "Representation learning: A review and new perspectives". In: IEEE Transactions on Pattern Analysis and Machine Intelligence 35.8 (2013), pp- 1798-1828. DOr:
10. 10.109/TPPMI 2013.5 50
[25]
Brenden M Lake et al. "Building machines that learn and think like people". In: Behavioral and Brain Sciences 40 (2017), e253, DOE: 10. 10110000000055211101111131111
[26]
t Jonas Peters ct al. Elements of causal inference: foundatioons and learning algorithms. The
MIT Press, 2017. URL:
causal inf erence
[27]
Yann LeCun et al. "Deep learning". In: Nature 521.7553 (2015), pp. 436-444. DOI: 10
1038/nature14539
[28]
Jürgen Schmidhuber. "Learning factorial codes by predictability minimization". In: Neural Computation 4.6 (1992). pp. 863-879, DOI: 10. w1162/necoo 1992. 4. 6.863
[29]
Michael Tschannen al. "Recent advances in autoencoder--assel representation learning".
In: (2018). arXiv: 1812. 05069 [cs.LG]
[30] Martin Schrimpf et al. "Brain-score: Which artificial neural network for object recognition is
most brain-like"" In: BioRxiv (2018), p. 407007. DOI: 10. 1101/407007
[31] Daniel LK Yamins et al. "Performance-optimize" hicrarchical models predict neural responses in higher visual cortex". In: Proceedings of the National Academy of Sciences 111.23 (2014),
pp. 8619-8624. DO1: 10. 1073/ 1073/pnas....// http3//nas.14433330. 1403112111
[32] Seyed- Seyed-Mahdi Khaligh-Razz Razavi and Nikolaus Kriegeskorte. "Deep supervised, but not unsu- pervised, models may explain IT cortical representation" In: PLoS Computational Biology 10.11 (2014), c1003915. DOI: 10. 1111//37oronaa pcbi pwpppccc........ 1003915
[33]
Daniel LK Yamins and James J DiCarlo. "Using goal-driven deep learning models under- stand sensory cortex". In: Nature Neuroscience 19.3 (2016), pp. 356-365. DOI: 10. 1038/nn nn
4244
[34]
Patrick Mineault al. "Your head is there move you around: Goal-driven models of the primate dorsal pathway". In: Advances in Neural Information Processing Systems. Ed. by M. Ranzato al. Vol. 34. Curran Associates, Inc., 2021, pp. 28757-28771. URL: https
/ //papers nips cc /wperpper////
Abstract.html
[35]
Eric Elmoznino and Michael Bonner. "High perforrinformingggg neural network models visual cortex benefit from high latent dimensionality". In: bioRxiv (2022), pp. 2022-07, DOI:
10.1101/2022 07 13. 439999999
[36] Colin Conwell et al. "What can 1.8 billion regressions tell us about the pressures shaping high-level visual representation in brains and machines?" In: bioRxiv (2023). DOI: 10. 1101/
2022.03
[37] Nicholas J Sexton and Bradley C Love. "Reassessing hierarchical correspondences be- tween brain and deep networks through direct interface". In: Science Advances 8.28 (2022), eabm2219, DOT: 10. 1126/ 1126/sciadv sciadv abm2219
12

---

[38] Greta Tuckute et al. "Many but not all deep neural network audio models capture brain responses and exhibit correspondence between model stages and brain regions". In: bioRxiv (2023). DOT: 10. 1101/2022. 09. 06. 506680
[39]
Blake Richards et al. "The application of artificial intelligence to biology and neuroscience".
In: Cell 185. (2022), pp. 2640-26443 DO1: 10. 1016/j cell. 2022. 06. 047
[40] Anthony Zador et al. "Catalyzing next-generation Artificial Intelligence through NeuroAl". In: Nature Communications 14.1 (2023), p. 1597. Dot: 10. 1038/
[41]
Adrien Doerig et al. "The neuroconnectionist research programme". In: Nature Reviews Neuroscience (2023), pp. 1-20. DO1:
[42] Nancy Kanwisher et al. "Using artificial neural networks to ask "why" questions of minds
and brains". In: Trends in Neurosciences (2023). DOI: D0.:[10.1]] 1016/ tins. 2022. 12. 008
[43] Rosa Cao and Daniel Yamins. "Explanatory models in neuroscience: Part I l-taking mechanis- tic abstraction seriously". In: (2021). arXiv: 2104. 01490v2 [i-
[44] Blake A Richards et al. "A deep learning framework for neuroscience". In: Nature Neuro-
science 22.11 (2019), pp. 1761-1770. DOI: 10. 00/3//415593-01-11
[45] David GT Barrett et al. "Analyzing biological and artificial neural networks: challenges with opportunities for synergy?" ?" In: Current Opinion in Neurobiology 55 (2019), pp. 55 64. DO1:
10. 1016/). 2019. 2019.01.007 01.007
[46]
Thomas Serre. "Deep learning: the good, the bad, and the ugly". In: Annual Review of Vision Science 5 (2019), pp. 399-426. DO1: 10. 1 1146/ annurev- vision -wwpaon--11---- 0917180014951
[47] Nikolaus Kriegeskorte. "Deep neural networks: a new framework for modeling biological vision and brain information processing". In: Annual Review of Vision Science (2015),
pp. 417-446, DOI: 10. 1101//29977
[48] James J Gibson. "The visual perception of objective motion and subjective movement". In: Psychological Review 61.5 (1954), p. 304, DO1: 10. 1037/h0061885
[49]
Leslie Ungerleider and Mortimer Mishkin. "Two cortical visual systems". In: Analysis of Visual Behavior (1982), pp. 149-586. URL: https: // www. cns nyu edu/ edu/"tony vns/
readings/ readinggs/ugerzeiderr wwwp://eeeeeeeeereeereree wwwidderrmishheeee mww.-miishiii-ie 1982. pdf
[50] Melvyn A Goodale and A David Milner. "Separate visual pathways for perception and action". "
In: Trends in Neurosciences 15.1 (1992), pp. 20-25. Do1: 10. 1ww.10066166-61 2236 (92))9449
[51] L. G. Ungerleider and L. Pessoa. "What and where pathways". In: Scholarpedia 3.11 (2008). revision #91940, p. 5342. DOr: 10. 01444444///aa ht////cholaaeeiia .3342
[52] Arash Vahdat and Jan Kautz. "NVAE: A Deep Hierarchical Variational Autoencoder".. In: Advances in Neural Information Processing Systems. Vol. 33. Curran Associates. Inc., 2020.
pp. 19661- 19679. URL: /
http://ppppp..lllaa///////////// html
[53] Mandyam Vecrambudi Srinivasan et al. "Predictive coding: fresh view of inhibition in the retina". " In: Proceedings of the Royal Society of London. Series B. Biological Sciences (1982), pp. 427-459. DOL: 10. (098/rrpp... 1982. 0085
[54]
Andre M Bastos et al. "Canonical microcircuits for predictive coding". In: Neuron 76.4 (2012), pp. 695-711. DOI: 10. 1tw6/:.neeronnn..... www.oe.roonoom 2012. 10. 038
[55]
Dawei W Dong and Joseph J Atick. "Temporal decorrelation: a theory of lagged and non- lagged responses in the lateral geniculate nucleus". In: Network: Computation in Neural Systems 6.2 (1995), p. 159. DOI: 10.
[56] Wolf Singer. "Recurrent dynamics in the cerebral cortex: Integration of sensory evidence with stored knowledge". In: Proceedings of the National Academy of Sciences 118 (2021). DO1: 10. 1ww0//oas..... 210104311
[57] Fabian A Mikulasch al. "Where is the error? Hierarchical predictive coding through dendritic error computation". In: Trends in Neurosciences 46.1 (2023), pp. 45-59. DOI:
10. 10.1016/j. tins 20222022222 .9..99.000 007
[58]
Beren Millidge al. "Predictive Coding: Towards Future of Deep Learning beyond Backpropagation?" In: International Joint Conference on Artificial Intelligence. 2022. DOI:
10. 24963/ijcai. 2022/774
[59] David C Knill and Whitman Richards. Perception as Bayesian inference, Cambridge Univer-
sity Press, 1996. Don: 10. N001101//6099780119984000
13

---

[60] Yair Weiss et al. "Motion illusions optimal percepts". In: Nature Neuroscience 5.6 (2002),
pp. 598-606 604.
[61]
Wilson Geisler and Daniel Kersten. "Illusions, perception and Bayes". In: Nature Neuro- science 5.6 (2002), pp. 508 510. DOL: 10. 1038/nn0602-503
[62] Iris Vilares and Konrad Kording. "Bayesian models: the structure of the world, uncertainty, behavior, and the brain".. In: Annals of the New York Academy of Sciences 1224.1 (2011),
pp. 22-39. D01: 10. 1111/j 1749-6632, 2011 05965.x
[63] Richard Langton Gregory. "Perceptions as hypotheses". In: Philosophical Transactions of the Royal Society of London. B, Biological Sciences 290.1038 (1980), pp. 181-197. DOT:
10. 10..09/9/STT. 1980. .0900
[64]
Timm Lochmann and Sophie Deneve. "Neural processing as causal inference". In: Current Opinion in Neurobiology 21.5 (2011), I), pp. 774-781. DOI: 10. 1016/006jjjj conb. 2011 05. 018 [65] Sabyasachi Shivkumar et "A probabillitiic population code based on neural samples", In: Advances in Neural Information Processing Systems. Ed. by S. Bengio et al. Vol. 31. Curran Associates, Inc., 2018. URL: https: /tttpps:/papprse nips cC/ wwwpcpppp..... pctppppr_fflleeeee_eeeeeeeeeeeee files paper/2018/
acf 2ee3/400011e63361111 7761336666677774042222677333333 Atw.-6666667AA-bbaaabcaaaa Abstract html
[66] József Fiser al. "Staaisticallyy optimal perception and learning: from behavior to neural representations". In: Trends in cognitive sciences 14.3 (2010), pp. 119-130. DOT: 10.
j.tics. 201cs.2010..0 01.
[67]
Bruno A. Olshausen. "Perception as Inference Problem". In: The Cognitive Neurosciences (5th edition) (2014). Ed. by Michael Gazzaniga and George R. Mangun. DOI: 10. 7551/ mitpress/ 9504. 003 0037
URL: attp://reet wrt/bruno////ooooo..otr..w..rrtt
bruno wttp:/ haapnn//pppr////a/e /pappptspeecoptioon
as-inference. p
[68]
Ferenc Csikor et al. "Top-down effeets in an early visual cortex inspired hierarchical Vari- ational Autoencoder". In: SVRHM 2022 Workshop @ NeurlPS. 2022. URL: https: //
openreview nwt. wwwriiewweetffrumm/ satpuummidddssfooooE0Y33
[69] Eleni Miliotou et al. "Generative Decoding of Visual Stimuli", In: Proceedings of the 40th International Conference on Machine Learning. Ed. by Andreas Krause et al. Vol. 202.
Proceedings of Machine Learning Research, PMLR, July 2023, pp. 24775-24784, URL: html [70] Yujia Huang et al. "Neural Networks with Recurrent Generative Feedback". In: Ad- vances Neural Information Processing Systems. Vol. 33. Curran Associates, Inc., 2020. URL: https / papers nips http://bbbbbbbett-a-baaa.......e
https: https://prooceeding: mlr. prwssrrsssseeee v2tss/v22//////2
paper files paper / 2020 hash /
[71] Rewon Child. "Very Deep (VAE]s Generalize Autoregressive Models and Can Outperform Them on Images" In: International Conference on Learning Representations. 2021. URL:
https://openrevieet htt::/oorommm///A-RAR
[72] Casper Kaae Sønderby et al. "Ladder Variational Autoencoders". In: Advances in Neural Infor- mation Processing Systems. Vol. 29, Curran Associates, Inc., 2016. URL: http:::ttt://ppeetet nips. coppaperr lemner_ffieee wwl.fiies/eleeeef http://ww...0.eee////////////ata
Abstract.html
[73] Lars Maaløe et al. "BIVA: A Very Deep Hierarchy of Latent Variables for Generative Model- ing". In: Advances in Neural Information Processing Systems. Vol. 32. Curran Associates, Inc., 2019. URL: hhtps: papers nips cc /paperppp.... files / paper / 2019 hash
[74] Louay Hazami et "Efficientvdvae: Less is more". In: (2022). arXiv: 2203. 13751v2 [cs.LG] LG]
[75] Shital Shah al. "Airsim: High-fidelity visual and physical simulation for autonomous vehicles". In: Field and Service Robotics: Results of the 11th International Conference.
Springer. 2018, pp. 621-635. DO1: 10. 1001 00007/978-3 1-3-319-777611--11
[76] Nicole C Rust and Anthony Movshon. "In praise of artifice". In: Nature Neuroscience 8.12
(2005), pp. 1647-1650. DOI: 10. 1038/ 1606
[77] Bela Julesz. "Foundations of cyclopean perception". In: (1971). URL: https //books google com
cttelle.com/ee www.coomoooseeee. /tt.:/oooooo/oooeeeoooo about wttooof__cooooe Cac__CCylleaanncmmmmammaa Periraatnneeppitieiooe html ?id=K
NfDgAACAAJ
14

---

[78] Tal Golan et al. "Controversial stimuli: Pitting neural networks against each other as models of human cognition". In: Proceedings of the National Academy of Sciences 117.47 (2020),
pp. 29330.29337. DOT: 10. 1073/pnas 1912334117
[79] Michael Beyeler al. "3D visual response properties of MSTd emerge from an efficient, sparse population code". . In: Journal of Neuroscience 36.32 (2016), pp. 1399-8415 DOI:
10. 15.35/2/JJT JNEUROSEU 009603966660 2016
[80] James J Gibson. "The perception of the visual world". In: (1950). URL: https://psycnet
apa.
[81] William H Warren Jr and Daniel J Hannon. "Direction of self-motion is perceived from
optical flow". In: Nature 336.6195 (1988), pp. 162-163. DOL: 10. 10w.10083333122
[82] J. Inigo Thomas et al. Spherical retinal flow for a fixating observer. Tech. rep. 1994. URL:
https:
2al7-a51fb7329588. upenn wewp/eetiiies//eeeeeeeeee htt.iiiieee/bbllaaaaaa//44444449 54cd-4833
[83] Jonathan Samir Matthis et al. "Retinal optic flow during natural locomotion". In: PLOS Computational Biology 18.2 (2022), e1009575. DOI: 10. 1371/journal
[84]
Eddy 1lg et al. "Flownet 2.0: Evolution of optical flow estimation with deep networks". In: Proceedings of the IEEE conference computer vision and pattern recognition. 2017, pp. 2462-2470. DOT: 10. 1009/ CVPR. 2017.
[85] Irina Higgins et al. Vbeta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework". In: International Conference on Learning Representations. 2017. URL: https
//openreview.oo
Sy2fzU9g1
[86]
Christopher P Burgess et al. "Understanding disentangling in 8- 8-VAE". In: (2018). arXiv: 1804. 03599 ML]
[87]
Nikolaus Kriegeskorte and Jör Diedrichsen. "Peeling the onion of brain representations", In:
Annual Review of Neuroscience 42 (2019), pp. 407-432. DOT: 10. 1144//nnuree--eurr
080317-061906
[88] Simon K Rushton and Paul A Warren. "Moving observers, relative retinal motion and the detection of object movement". In: Current Biology 15.14 (2005), R542-R543, DOI: 10. 1016/0...ccu cub. 2005 .200
[89]
Paul A Warren and Simon K Rushton. "Optic flow processing for the assessment of object movement during ego movement". In: Current Biology 19.18 (2009), pp. 1555-1560. DOI:
10. 1016/j cub. 2009000000 07.000.07 057
[90]
Paul A Warren and Simon Rushton. "Perception of object trajectory: Parsing retinal motion into self and object movement components". In: Journal of Vision 7.11 (2007), pp. 2-2. DOT:
10. 10.1167/7.11.
[91]
Nicole E Peltier et al. "Optic flow parsing in the macaque monkey". In: Journal of Vision
20.10 (2020), pp. 8-8. DOI: 10. 1167/jov. 20. 10.8
[92] James C. R. Whittington al. "Disentanglement with Biological Constraints: Theory of Functional Cell Types" In: The Eleventh International Conference on Learning Representa- tions. 2023. URL: https: /ttp://////eeeeieee// net/forumiid-990 Gf
[93] Sebastien Lachapelle et al. "Synergies between Disentanglement and Sparsity: General- ization and Identifiability in Multi-Task Learning". In: Proceedings of the 40th Interna- tional Conference on Machine Learning. Ed. by Andreas Krause et al. Vol. 202. Proceed- ings of Machine Learning Research. PMLR, July 2023, pp. 18171-18206, URL: https html
//proceedings mlr.
[94] Abhishek Kumar et al. "Variational Inference of Disentangled Latent Concepts from Unla- beled Observations" " In: International Conference on Learning Representations, 2018. URL: net/ wwwoorrmmddiiiioo
https://openreva/
[95] Hyunjik Kim and Andriy Mnih. "Disentangling by factorising", In: International Conference on Machine Learning. PMLR. 2018 pp. 2649-2658. URL: httt:::/proceediiii.. mlr Peess/v80/king
html
[96] Ricky T. Q. Chen et al. "Isolating Sources of Disentanglement in Variational Autoencoders". In: Advances in Neural Information Processing Systems. Ed. by S. Bengio et al. Vol. 31. Curran Associates, Inc., 2018. URL: http:: nips. /ts.pp..c/pappr.ilee///eeeeeeeee
15

---

[97] Michal Rolinek et al. "Variational autoencoders pursue PCA directions (by accident)". In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2019, pp. 12406-12444 12445. URL:: https: openaccesss thecvf com/ content CVPR (049/htm//Rolle wetooline_Varaatiiiio Variationall Aut.oeecoeeer Accident_
htt.:///eoeoooeiaeeeaaaaaaaaaaar Sttpooedddrrruuuuuueeeee Pursue_ SOAA_Directions by
pap.rr.pp..... html
[98] Sjoerd van Steenkiste al. "Are Disentangled Representations Helpful for Abstract Visual Reasoning?"" In: Advances in Neural Information Processing Systems. Vol. 32. Curran Asso-
ciates, Inc., 2019. URL: https: /tttpp:://pppprs...m www..iiipcc........ httpc///ppper.fleeeppappapppepee
bc3c44a6331$8995504444 -aw.:9950455aaaaaababb-------- a wwt.oaaaaa-stttttttttmmm html
[99] Andrea Dittadi et al. "On the Transfer of Disentangled Representations in Realistic Set- tings". In: International Conference on Learning Representations. 2021. URL: https
//openreview. net/ forum? dd-VVvvjiivjjooooo 1QNR11
[100] W Jeffrey Johnston and Stefano Fusi. " "Abstract representations emerge naturally in neural networks trained to perform multiple tasks". In: Nature Communications 14,1 (2023), p. 1040.
DOI: 10. 1038/ 0008://s14477-000 023 36583-0
[101] Mare-Andrt Carbonneau al. "Measuring disentanglement: A review of metrics". " In: IEEE Transactions on Neural Networks and Learming Systems (2022). DO1: 10. 1109/TNNLS TNNLS
2022. 3218982
[102]
Cian Eastwood et al. "DCI-ES: An Extended Disentanglement Framework with Connections to Identifiability". In: The Eleventh International Conference Learning Representations. 2023. URL: https: http///pperrvveeeeeee net/ffrunf ftr.moormmmi/ idd-22z 4622-ggghhttt
[103] Yuwei Cui et al. "Diverse suppressive influences in area MT and selectivity to complex motion features". In: Journal of Neuroscience 33,42 (2013), pp. 16715-16728. DO1: 10
1523/JNEUROSCI. SBLEUUROCII.200-1111 0203- 2013
[104] Yuwei Cui et al. "Inferring cortical variability from local field potentials". In: Journal of Neuroscience 36. (2016), pp. 4121-4135, DOI: 10. J523/ONEUOOSCI 2502- 15. 2016 [105] Maneesh Sahani and Jennifer Linden. "How Linear are Auditory Cortical Responses?" In: Advances in Neural Information Processing Systems. Vol. 15. MIT Press, 2002. URL: https / papers nips cc paper - files paper / 2002 / hash / 7b4773 (18817782888888811111111- Abttrattt..--aatttt........
[106] Patrick Mineault al. "Hierarchical processing of complex motion along the primate dorsal visual pathway". . In: Proceedings of the National Academy of Sciences 109.16 (2012),
E982-E980. DOI: 10. 1073/pnas 1115685109/
[107] Eero P. Simoncelli and Bruno A. Olshausen. "Natural image statistics and neural representa-
tion." " In: Annual Review of Neuroscience 24 (2001), pp. 1193-216. DOI: 10. 1146/ annurev
neuro 24. 1. 1193
[108] Aaron van den Oord et al. "Representation Learning with Contrastive Predictive Coding". In: (2019). arXiv: 1807 03748 [cs LG]
[109]
Karl S Muller et al. "Retinal motion statistics during natural locomotion". In: eLife 12 (2023),
e82410. DOI: 10. 7554/eLife 82410
[110] Benedict Wild and Stefan Treue "Primate extrastriate cortical area MST: gateway between sensation and cognition". In: Journal of Neurophysiology 125.5 (2021), pp. 1851-1882, DOI:
10. 10.1152/jn. 00384. 2020
[111] Benedict Wild et al. "Eleecrophysiological dataset from macaque visual cortical area MST in response to a novel motion stimulus". In: Scientific Data 9 (2022). DOT: 10. . 4038/ss1597)
022- 02239923--
[112] Alex H Williams et al. "Generalized Shape Metrics on Neural Representations", In: Advances in Neural Information Processing Systems. Vol. 34. Curran Associates, Inc., 2021. URL: https: // papers nips cc paper files / paperr / 2021 hash /
http:aaaaebbbeee6222222222233333 -aataaaaeeeeeebbaaaaaaa.....a... www.tbast.tttb..m html
[113] Lyndon Duong al. "Representational Dissimilarity Metric Spaces for Stochastic Neural Networks". In: The Eleventh International Conference on Learning Representations. 2023. URL: https:://ppenreiiieet
[114] Max Klabunde et al. "Similarity Neural Network Models: A Survey of Functional and Representational Measures". In: (2023). arXiv: 2305. 06329 LC]
16

---

[115]
Abdulkadir Canatar et al. A Spectral Theory of Neural Prediction and Alignment. 2023. arXiv: 2309. 12821 [q- Ag-bio.NO.
[116]
Adam Paszke et al. "PyTorch: An Imperative Style, High- Higgh-Performacce Deep Learning Library". In: Advances in Neural Information Processing Systems. Vol. 32. Curran Associates, papers nips cc pwpprrappprrr__ files paper/ppee///// / 2019 hash / Inc., 2019. URL: https: /
Abstaat hhtml
[117] Charles R. Harris al. "Array programming with NumPy". In: Nature 585.7825 (Sept. 2020).
pp. 357-362. DO1: 10. 10080008//0
[118] Pauli Virtanen et al. "SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python".
In: Nature Methods 17 (2020), pp. 261-272. DO1: [0. 1038/ 15244-922220 00099-066-222 2
[119] Fabian Pedregosa et al. "SScikitlearn: Machine learning in Python". " In the Journal of Machine
Learning Research 12 (201 1), pp. 2825-2830. DOI: 10. 5555/5555555 1 1953048 2078195
[120] The pandas development team. pandas- devndad-dv//aaddss: Pandas. Version latest. Feb. 2020. DOE:
10. 00.521/zzeoodo. 3509134
[121] John D Hunter. "Matplotlib: A 2D graphics environment". In: Computing in Science & Engineering 9.03 (2007), pp. 90-95, DOI: 10. 1109/MCSEE 2007
[122] Michael L Waskom. "Seaborn: statistical data visualization" In: Journal of Open Source
Software 6.60 (2021), p. 3021. DO1: 10. 200.2115//o88... 03021
[123] Edward H Adelson and James R Bergen. "Spatiotemporal energy models for the perception
of motion".. In: Josa a 2.2 (1985), pp. 284-299. DOr: 10. 10011334/J00 JOSAA. 2. 000284
[124] Shinji Nishimoto and Jack L Gallant. "A three-dimensiones spatiotemporal receptive field model explains responses of area MT neurons to naturalistic movies". In: Journal of Neuro- science 31.41 (2011), pp. 14551-14564. DOI: 10. 1523 JNEUROSCUR 6801-8601111 2011
[125] Yena Han et al. "System identification of neural systems: If we got right, would we know?" In: International Conference on Machine Learning. PMLR. 2023, pp. 12430-12444. URL.: https: httpass///ccceeengs...mm mlr. watpeessvv222/aaae w22///a222d...... .ttml
[126] Jie Hu et al. "Squeeze -ndzzzz-------ae and exciittttttioon networks". In: Proceedings of the IEEE Conference Computer Vision and Pattern Recognition. 2018, pp. 7132-7141. DO1: 10. 1109/ CVPR
2018.0011 00745)
[127] Sergey loffe and Christian Szegedy. "Batch normalization: Accelerating deep network training by reducing internal covariate shift". In: International Conference on Machine Learning, pmlr. URL: hhttps://ppoceeeeeee
2015, pp.
/vassss/v///// mlr. prwpppressseee vv7 1//3//offfa /1offffffff.////////////fffoffff html
[128] Tim Salimans and Durk P Kingma. "Weight Normalization: A Simple Reparameteri- zation to Accelerate Training of Deep Neural Networks". In: Advances in Neural In- formation Processing Systems. Ed. by D. Lee et al. Vol. 29. Curran Associates, Inc., 2016. URL: https: / papers nips cc / paper files / paper / 2016 / hash / Abstreetbbbtttc.c.... html [129] Prajit Ramachandran et "Searching for Activation Functions", In: International Confer- ence on Learning Representations. 2018. URL: htttps://openreviiee net/forum?? idd=
SkBYYyZRZ
[130] Stefan Elfwing et al. "Sigmoid-weighted linear units for neural network function approxima- tion in reinforcement learning". In: Neural Networks 107 (2018), pp. 3-11 11. DOI: 10 1016/]
neunet.2017. 12.02..000 012)
[131] Yuichi Yoshida and Takeru Miyato. "Spectral Norm Regularization for Improving the Gener- alizability of Deep Learning". In: (2017). arXiv: 1705. 10941 [stat. ML.]
[132] Samuel R. Bowman et al. "Generating Sentences from a Continuous Space". In: Proceedings of the 20th SIGNLL Conference on Computational Natural Language Learning. Berlin, Germany: Association for Computational Linguistics, Aug. 2016, pp. 10-21. DOL: 10. 18653/ v1/K16-1002
[133] Hao Fu et al. "Cyclical Annealing Schedule: A Simple Approach Mitigating KI. Vanishing". In: Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers). Minneapolis, Minnesota: Association for Computational Linguistics, June 2019,
pp. 240-250. DOL: 10. 1ww.1/81///1111-1 1021
17

---

[134] Arash Vahdat et al. "DVAE++: Discrete Variational Autoencoders with Overlapping Transfor- mations". In: Proceedings of the 35th International Conference on Machine Learning. Ed. by Jennifer Dy and Andreas Krause. Vol. 80. Proceedings of Machine I Leannigg Research. PMIR. July 2018, pp. 5035-5044. URT.: wtw.:tttt://prrccddingeeee http://iinwni.....eeeeeeeeeeeeee 2a.8000vaaat1.... html
mlr
[135] Xi Chen al. "Variational Lossy Autoencoder".. In: Interatioonal Conference on Learning Representations. 2017. URL: hattps://oppereeeeeetee
[136]
Diederik P Kingma and Jimmy Ba. "Adam: A method for stochastic optimization". In: (2014).
Xrx: 1412. 6980 [cs LC]]
[137] Ilya Loshchilov and Frank Hutter. "SGDR: Stochastic Gradient Descent with Warm Restarts". In: International Conference on Learning Representations. 2017. URL: https
//openreview net/ ctnttfforumiidde ix
[138] Yoav Benjamini and Yosef Hochberg. "Controlling the false discovery rate: a practical and powerful approach to multiple testing", In: Journal of the Royal Statistical Society; series
(Methodological) 57.1 (1995), pp. 289-300. DO1: 10. 1.0.111//,... 1111111.2117611111 1995. TB02031
[139] Jacob Cohen. Statistical power analysis for the behavioral sciences. Academic press, 1988. DOL: 10. #307/2529115
[140] Gregory C. DeAngelis and Dora F. Angelaki. "Visual-Vestibule Integration for Self-Motion Perception". In: The Neural Bases of Multisensory Processes (2012), pp. 629-644 DOI:
10. 1201/9781 439812174
[141]
Eduard Von Holst. "Relations between the central nervous system and the peripheral organs". In: British Journal of Animal Behaviour (1954). DOT: 10. 10166/099005600 (54) 80044-X
[142] Paul R MacNeilage al. "Vestibular facilitation of optic flow parsing". In: PLoS One 7.7 (2012), e40264. DOE: 10. 1 1 journal. pone. 0040264
[143] Kathleen E Cullen and Omid A Zobeiri. "Proprioception and the predictive sensing of active self-motion", In: Current Opinion in Physiology 20 (2021), pp. 29-38. Do1: 10. 1016/j
cophys. 2020. 001
[144] Constance S Royden and Ellen C Hildreth. "Human heading judgments in the presence of moving objects". In: Perception & Psychophysics 58 (1996). pp. 836-856. DOI: /7557
BF03205487
[145] William H Warren Jr and Jeffrey A Saunders. "Perceiving heading in the presence of moving objects". In: Perception 24.3 (1995), pp. 315-331. DOI: 101:10.:1 1a::]10111600 1001/022003311
[146] Edward AB Horrocks et al. "Walking humans and running mice: perception and neural encoding of optic flow during self-motion", In: Philosophical Transactions of the Royal
Society B 378. 18.1166 869 (2023), 20210450. DOr: 10. 1098/rstb. 2021. .4450
[147] Jean-Pau Noel al. "Causal inference during closed-loop navigation: parsing of self-and object-motion" " In: Philosophical Transactions of the Royal Sociery B 178.1886 (2023), 2022.022222 0344
p. 20220344. DOT: 10. 1ew.1099/rst
[148] Denis N Lee. "The optic flow field: The foundation of vision". In: Philosophical Transactions
of the Royal Society of London. B, Biological Sciences 290.1038 (1980), pp. 169-179. DOI:
10. D0.1098/rsst..o 1980.9990. 0089
[149] Markus Lappe et al. "Perception of self-motion from visual flow". " In: Trends in Cognitive Sciences 3.9 (1999), pp. 329-336. DO1: 10. 1016//1364----1 6613 (99) 01364-9
[150] Irina Higgins al. "Symmetry-based" representations for artificial and biological general intelligence". In: Frontiers in Computational Neuroscience 16 (2022), p. 836498. DOr:
10. ww..39ffffoom.ommo 2022.836498
[151]
Fabio Anselmi et al. "On invariance and selectivity in representation learning" In: Information and Inference: A Journal of the IMA 5.2 (2016), pp. 134-158. DOI: 10. 1093/ imaiai/ iaw009]
[152] Michael M Bronstein et al. "Geometric deep learning: Grids, groups, graphs, geodesics, and gauges". In: (2021). arvi:: 2104. 13478 [cs. LC]
[153] Ishaan Gulrajani al. "PixeIVAE: A Latent Variable Model for Natural Images". In: Interna- tional Conference on Learning Representations, 2017. URL: https: /tttps://openreviees net/ Forum?id-BJJVvvt51g
18

---

not necessarily Supplementary material for:
Hierarchical VAEs provide a normative account of motion processing in the primate brain
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 124]. This strategy also allows for direct comparison with recorded neural data in MT and MST using random dot kinematograms [103] 104 [06]. However, a more complex model would be necessary to explain responses of neurons earlier in the visual pathway, such as V1, which would require pixel-computable model as in previous work (e.g., DorsalNet (34).. Likewise, our > 2x 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 equivalent to the random dot velocity fields used in the experiments. Thus, it is an open question whether hierarchical VAE trained and tested on video stimuli would better align 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 fair comparison.
123
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 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), expect that hierarchical computation is 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 aligmment that is sensitive to the sparsity of latent-to-neuron relationships (" ("aligmment-score" Fig, g 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 [T25] measures of representational alignment. Consistent with this, our application of linear regression found that 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 aligmment score (Fig. (a) was much more consistent distinguishing between models (see Fig. [6], and revealed that hierarchical models (both cNVAE and cNAE) had significantly sparser latent-to-neuron relationships. The alignment score measures whether model has learned similar representational "form" " to the brain, which would enable the sparsity of relationships. This concept closely related the "completeness" score from the DCI framework (22). The alignment score shown in Fig. & 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 proper metric in the mathematical sense [112). Future work will consider more sophisticated metries for brain alignment
112
19

---

Additional Methods
9.1
VAE loss derivation
Suppose some observed data are sampled from generative process as follows:
p(a:)= = p(ar, z)dz
f p(zrl))p(z)ds dz,
(1)
where z are latent (or unobserved) variables. In this setting, is interesting to ask which set of latents are likely, given an observation æ. In other words, we are interested in posterior inference
p==(2) x p(z)=))(z).
The goal VAΣS to approximate the (unknown) true posterior p(z (x)) with a distribution q( g(=|a; ; 0), where 0 are some free parameters to be learned. This goal is achieved by minimizing the Kullback- Leibler divergence between the approximate posterior and the true posterior p:
Goal:
minimize
DEL [<</==;)) 0) p(=]
This objective is intractable. but rearranging the terms leads to the following loss that is also the (negative) variational lower bound on log p(z):
LVVE = -F.g log p(x)p( 0awe) DKL. [(( ax; Oene)
p(= (0
(4)
where Oene and Odeo are the parameters for the encoder and decoder neural networks (see Fig. 2).
The first term in [Equation 4|is the reconstruction loss, which chose 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 VAFs, the prior not parameterized. Instead, is chosen to be simple distribution such Gaussian with zero mean and unit covariance:
Odcc) p(z) = I).
(5)
The approximate posterior is also a Gaussian with mean pa( Oene) and variance o² (æ;; O Oene):
<( az; Oene) = (Ns( Oene), (x;; Oene))
(6)
As result, we see that the KL term for vanilla VAE only depends on encoder parameters 0enc-
9.1.2 Prior, approximate posterior, and the KL term in the eNVAE
Similar to the NVAE (5), 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.c., 21 to Ø3 in Fig. 2h. 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" indicated as 22 and 23 respectively: the cNVAE is designed in way that changing the parameters along the decoder pathway will impact the prior distributions on 22 and 2y (but not 21). Note the "h" in Fig. 2] which is also learnable parameter. In summary, the "top" cNVAE latents (e.g., 21 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.
latent groups in Fig. 2
More formally, the eNVAE latents are partitioned into disjoint groups, z = {21, #2, z1.) where L is the number of groups. Then, the prior is represented by:
20

---

pl=)0 - p(z1)
II t=2
(z 265 Ocee)s
where the conditionals represented by factorial Normal distributions. For the first latent group have p(z1) = (0, I). similar to vanilla VAEs. For every other latent group. we have:
dec) = (u(zze) ddec), (z ddec));
(8)
where p(z<e; dee) and a² (=<t: e:: dee) are outputted from the decoder sampler layers. Similarly, the approximate posterior in the cNVAE is represented by:
L II E=1
(22/2; Oenc)
(( Zel, az Oenc).
(9)
We adopt a Gaussian parameterization for each conditional in the approximate posterior:
q(ze x. Onne) -NW Oenc), o² o² (z x; Oene)),
(10)
where p1=ce,, Deno) and o (= I;; O eno) are outputted from the encoder sampler layers (Fig. grey trapezoids). We are now in a position to explicitly write down the KL term from Equation for the cNVAE:
2
1. Σ
KL term DKI. |(21|2e, Oene) p(z1)
+
KLe(Cenoc 0amr))
(11)
where KL refers to the local KL term for group and is given by:
KLe(Oeno) enc. dsec) DKL (g (ze acee x; Oene) (zelzzze, Qdec)
(12)
and the approximate posterior up to the (E 1)th group is defined as:
a(=<e x; 0 Onn))
II i=1
g( = Penc).
(13)
This concludes our derivation of the KI. terms shown in Table
9.2 Key architectural differences between NVAE and eNVAE
Our main contribution architecture design lies in the modification of the NVAE "sampler" layer and the introduction of the "expand" layer, which are simple deconvolution (Fig. 2 In the NVAE, the
purpose of sampler layer is to map encoder features to mean and variance vectors u(xx) and o² (x). which are then used to construct approximate posterior distributions y (z|æz ((((()), o² (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 inference. the NVAE, all sampler layers are convolutional with fixed kernel size of 3 x 3 and padding of 1. Thus, the application of a sampler layer does not alter the spatial dimension hidden encoder features, resulting in a convolutional latent space.
(²
For example, at level é of the hierarchy. μ (x) and of (æ) will have shapes like d X St X sz. where is number of latent variables per latent group, and 8, is the spatial scale of the processed encoder features at level f. As result, NVAE latents themselves end up being convolutional in nature, with shapes like 88 x 8;- This design choice leads to the massive expansion of NVAE latent dimensionality, which is often orders of magnitude larger than the input dimensionality.
21

---

To address this, we modified NVAE sampler layers in 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 eNVAE sampler layers such that their kernel sizes are selected adaptively to match the spatial scale of their input feature. For example, if latent group l operates at the scale of 8g, then the sampler layer will have kernel size of St X 8g, effectively integrating over space before the sampling step. This results in latents with shape d 1 x 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 kemel size of sy X Sg. Thus, processing d x 1 x 1 latent by an expand layer results in d X 8g X sy output, allowing them be combined with the decoder features (concatenation or addition; blue "+" in Fig.
cNVAE pros and cons. Our approach resulted much smaller latent space compared 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 uapu-size-dependerr 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 x 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 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 [28] 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 weal-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 hierarchy of latent groups that sampled sequentially: "top" latents are sampled first, all the way down to "bottom" latents that are closest 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 total of 21 hierarchical latent groups: 3 groups operating at the spatial scale of 2 x 2, 6 groups at the scale of 4 x 4, and 12 groups at the scale of x 8 (see Figs. 2 and 3 Each latent group has 20 latents, which yields 21 x 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 (7I). 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 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 otherwise identical but non-compressed NVAE. The total number of latent variables for hierarchical VAE is given by:
L
D
de
(14)
22

---

Table 4: The hierarchical models (cNVAE, cNAE) process information at various spatial scales. These models have total of 21 latent groups, each containing d = 20 latent variables. For hierarchical models, we concatenate the latent representations across levels of the hierarchy, resulting in single latent vector with dimensionality D 420. In contrast, non-hierarcchcal models (VAE, AE) do not possess such a multi-scale organization. Instead, their non-hierarchical latent spaces contain a single latent group, operating at single spatial scale. This single latent group contains the entire set of d = D = 420 latent variables. See Table 2 and Fig.
Number of latent groups
Number of latent variables per group
Latent space dimensionality
Model
2 x 2
4 x 4
8 x 8
total
cNVAE
3
6
12
21
d = 20
D = 21 20 420
VAE
1
1
= 420
= x 420 = 420
cNAE
3
6
12
21
20
21 20 420
AE
1
1
d = 420
D = x 420 = 420
where de is the number of latents in level É, and is the total number of hierarchical levels. As described above, de d sî for NVAE, where is the number of latents per level and is the spatial scale of level In contrast, this number a constant for the cNVAF: de = d. In this paper, we set d 20 and had the following configuration for the eNVAE:
"top" latents: 3 groups operating at the spatial scale of X 2
"mid" latents: 6 groups operating at the spatial scale of
"bottom" latents: 12 groups operating at the spatial scale of 8 x 8
This resulted in an overall latent dimensionality of D (3 + 6 + 12) x d 21 20 420 for the eNVAE. In contrast, an NVAE with a similar number hierarchical levels would have a dimensionality of D (3 2 + 6
42 + 12 8²) x d 876 20 520.
Note that Fig. 4 we necessarily had to reduce the NVAE depth to roughly match its latent dimensionality other models. Specifically, the NVAE in Fig. had 2/3/6 latent groups at top/middbottomm with d = 1. This configuration resulted in latent dimensionality of = (2 2 + 3 4a + 6 82) 440, which is roughly comparable to other models D 420. See Table for comparison between eNVAE and the alternative models used in this study.
9.4 Model training details
We generated 750, 000 samples for each ROFL category, with 600, 000/75, 000/75, 000 split for train/validationnnesst datasets. Model performance (reconstruction loss, NELBO) was evaluated using the validation set (Fig. [nd 15). Similarly, we used the validation set to estimate mutual information between latents, 2, 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 We then test the performance
using the test set. This includes results presented in Fig. 4 and Fig.
KL annealing and balancing. We annealed the KL term during the first half of the training, which is known be an effective trick in training VAΣS (52] 72 132 133). Additionally, during
the annealing period, we employed a KL balancing mechanism which ensures an equal amount of information encoded in each hierarchical latent group [22 134
[35).
Gradient norm clipping. During training, we occasionally encountered updates with large gradient norms. Similar to vdvae (7I], we used gradient clipping, with 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 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 cosine learning rate schedule
I37 but without warm restarts.
23

---

9.5 Choosing 3 values for different figures
known that different choices of the 3 value in the VAE loss (Table 2 will result in different model properties [85]. Therefore, while we scanned across many /3 values spanning two orders of magnitude, we displayed results for certain betas best demonstrate our points.
To ensure a fair model comparison in the results shown in Figs. 3
and 12 we selected / values that maximized the overall untangling score for each architecture (cN (cNVAE, $
0.15; VAF, 8 1.5). This can be seen in Fig. [ which presents DCI scores over the entire 3 spectrum. The top row of Fig. displays informativeness untangling) scores, revealing an inverted U shape, consistent with previous work 885 Different architectures peak at different / values.
For the neuron shown in Fig. 8b, found that different / values lead qualitatively similar outcomes, that is, extreme sparsity in eNVAE feature importances, in sharp contrast to that of VAE. We deliberately chose to show a larger 3 5 for VAE because previous work by Higgins et al. [16] suggested that increasing 3 values increases neuronal alignment, which we did not observe here. In contrast, even very small 0.01 for the eNVAE results in 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 3 values alone althohal we do observe some dependence B, so ultimately a combination of both architecture and loss is important (but mostly architecture).
9.6 Details on latent code evaluation
Each ROFL sample 2 is rendered from low dimensional ground truth vector g E R where 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 curved K -ddimensional manifold embedded in a 222-dimension 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 pixel at the exact center, which served the center of gaze or the fixation point (Fig. Th). The true dimensionality (number of generative factors) of different ROFL categories given in Table I For example, is K 11 for fixate-1.
We trained unsupervised models by feeding them raw samples and asking them to produce reconstruction é. The models accomplish this by mapping their inputs to latent codes 2 during inference (red pathways in Fig. 2). and mapping z to during generation (blue pathways Fig. 2 Note that for all models we have z E IR 420 (see Table 4h.
In sum, our approach consists of four steps:
1. Data generation:
ROFL +
2. Unsupervised training: + Model
+ +
3. Inference: a Model +
4. Evaluation: z + g?
Within this framework, we investigated the relationships between data generative factors 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 unsupervised models (Table and evaluated them based on (1) their ability to represent the generative variables of the ROFI. stimuli; and (2) the neural alignment of their representation. This resulted in total of five different metrics used evaluate model performance. Below, we describe the different metrics used for each.
#1: untangling disentangling of data generative factors. 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 comprised of three different metries, Disentanglement, Completeness, and Informativeness, which we will discuss in detail below.

---

Entangled
Overcomplete
One-to-one
91
92
93
92
91
92
93
23
2
22
23
24
25
21
2
33
24
25
Figure 9: Suppose we train a regression model predict ground truth factors 9j (blue) from latents (grey). Entanglement (left) refers to situation in which one latent variable contains information about many generative factors. overcomplete code (middle) 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 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. Sa.
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 good latent code idcally disentangles as many factors of variation as possible while discarding as little information possible. In other words, we consider latent code to be "good" long as:
1. It can predict the variance of data generative factors, and
2. It exhibits one-to-one relationship with the generative factors (Fig. 9
To make this idea concrete, the starting point is to train a regressor f (e.g., lasso) that predicts given data generative factor 9j from latent code z. The performance of f in predicting (measured using e.g., R22) is referred to as the Informativeness score. The notion of Informativeness, if measured using a linear regressor, is very similar Untangling (20].
As the next step, the weights of f are used to construct a matrix of relative importances R, where Rij denotes the relative importance of z; in predicting Jj- For example, be the magnitude of lasso coefficients. The matrix of relative importance is then used to define two pseudo-probability distributions:
Ry
Ps.
probability of z; being important for predicting 9j
(15)
Σn Rik
Rij
P. Σn Rkj
probability of gj being predicted by 2;
(16)
These distributions are then used to define the following metrics, which complete the DCI trinity:
Disenlanglement:
D, = 1 where Hx(Pr.) -E Pij log P.j
(17)
Completeness:
1 where HD(P,)) = Σ P.y log Prj
(18)
Here, K is the number of data generative factors (e.g., K = 11 for fixate-1, Table 1. and D is the number latent variables (D = 420 for the models considered in our paper, Table
Intuitively, if latent variable z; contributes to predicting a single ground truth factor, then D, = 1. Conversely, D. when 2; equally contributes to the prediction of all ground truth factors. the ground truth factor g, is being predicted from single latent variable, then C 1 (a complete code).
25

---

On the other hand. if all latents contribute to predicting 9j. then C 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 random forest as their choice of regressors f. Initially, we experimented with lasso but found that the lasso coefficient had strong impact on the resulting DCI scores. To mitigate this, used linear regression instead and estimated the matrix of relative importance using scikit-learns's 's permutation importance score (sklearn. inspantionnn permutation importance). which measures how much performance drops based on the shuffling of given feature.
9.7.2 Brain alignment score
As depicted in Fig. Sh, the alignment score emphasizes the degree of correspondence between latent variables and neurons. sparse relationship between latents and neurons suggests a higher alignment, indicating that the model's 's representation closely mirrors the neuron's 's representational "form" [16).
This notion of alignment is intimately tied the "completeness" score discussed above. See Fig.
Dense
Sparse
MT
MT
Z1
Z2
Z3
Z4 4
Z5
Z1
Z2
Z3
Z4
Z5
Figure 10: Consider training a regressor to predict MT neuron responses (green) from latent variables 2 (grey). Beyond accurate predictions, we assess the alignment between and MT based on the sparsity the resulting aatent-to-neuron relationships. One possibility that many latents contribute to predicting the neuron in dense fashion. This would suggest that the neuron's '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 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 hicrarchical levels (see section 9.3 for details). To demonstrate the difficulty of our untangling task, included two simple baselines: Raw data and 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 (~ 55%) and object velocity in and directions (47%). This result demonstrates the importance of nonlinear computation in extracting good representations.
10 Additional Results
In this section, we present additional results demonstrating the advantages of hierarchical VAEs versus con-hierarchical ones, which support the results in the main text. All models explored in this section were trained on fixate-1. For both eNVAE and VAE, we chose 3 values that maximized their informativeness (we discuss this choice in section 9.5 Specifically, we set B = 0.15 for cNVAE and 3 = 1.5 for VAE, These same models trained using the respective 3 values were also used for results presented in Fig. 3 and Fig. in the main text.
26

---

10.1 Individual cNVAE latents exhibit greater linear correlation with ground truth factors
In the main text (Fig- 4 we demonstrated that the cNVAE latent representations z contained explicit information about ground truth factors g. We achieved this by training linear regression models that map the 420-dimensional latent space to each ground truth factor. Here, we investigate the degree that individual latents 2; are (Pearson) correlated with individual factors 95 (Fig. [TI. For some ground truth factors, we find that cNVAE learns latents that are almost perfectly correlated with them. For instance,
0.91 for F2 and 0.87 for Fy. compared to modest correlation values of 0.34 and r = 0.38 for VAE. The higher linear correlation consistently observed for cNVAE over VAE. Especially, note that the VAE does not capture object-related variables all, consistent with other results in the main paper (Figs. [nd
For the CNVAE, the indices for the selected highly correlated latents are follows: (296, 248, 393, 284, 368, 92, 206, 105, 338, 60, 63], ordered left-to-right in the same manner as in Fig. II Here, index of 0 corresponds to the very top latent operating at the spatial scale of 2 2, and an index of 419 corresponds to the very bottom one at scale 8 x (see Fig. 2 By considering the hierarchical grouping of these latents (see dashed lines in Fig. 3 we see that the majority of the highly correlated latents operate at the spatial scale of 8, with the exception of object-related ones corresponding to Xobj. Zobj. Vobj.y. and Vobj.= which operate at x 4.
In conclusion, Fig. TT revealed that eNVAE learns individual latents that are highly correlated with individual ground truth factors, which provides complementary insights to why cNVAE is so successful untangling (Fig.4) However, we note that a high correlation with individual latents not the best way to capture the overall functional organization, which is better reflected for example in the mutual information heatmap of Fig.
Fy
Vunt
Veniy
Veni
Xay
Yay
Zoly
Vay
Vany CA
Vatj.
0.37
3.0
0.4
COVAE
tm
3:00
2:06
p.0
4.0
VAE
Figure 11: For each ground truth factor and model, we identified single latent variable that displayed the highest (Pearson) correlation with that variable. The ground truth factors exhibited greater linearity within the latent space of cNVAE compared to VAE. Scatter points correspond to data samples from the test set. x-axis, the value of ground truth factors; y-axis, the value of latent variables.
10.2 Generated samples: non-hierarchical VAE completely neglects objects
If generative model has truly learned the underlying structure of dataset, then it should be able to generate fake samples that look like its training data. Here, we examine generated samples from both the cNVAE and VAE models, as illustrated in Fig. [2] For comparison, we also include five samples randomly drawn from the fixate-1 dataset that these models were trained on (Fig. I2 first row). The vNVAE samples appear indistinguishable from true samples; however, VAE fails to capture and generate objects, corroborating our previous results (see Fig. 3 Fig. 4 and Fig. TT]. This striking divergence between eNVAE and VAE indicates that a hierarchically organized latent space necessary for effectively modeling datasets that contain multiple interacting spatial scales.
10.3
Latent traversal reveals multi-scale understanding in eNVAE
One desired property of VAE latent representations is that continuous interpolation along latent dimension results in interpretable effects on the corresponding reconstruction. Here, we provide visual examples of latent traversal in cNVAE, focusing on a latent dimension that effectively captures Yonj (Figure l3). These examples demonstrate intriguing interactions between object velocity patterns and self-motion within the scene. Notably, in some examples such as rows 0 and 6, relocating the object from the bottom to the top the scene induces a significant change its linear velocity along the y-axis. This phenomenon arises due to the influence of self-motion, seen in the background.
27

---

fixate-1
cNVAE (bierarchical)
VAE (noin www-hherrhhhhcall
Figure 12: The first shows random samples drawn from ROFL fixate-1 category. The second and third rows are samples generated from cNVAE and VAE models that were trained on fixate-1. Generated samples from cNVAE closely resemble real data; whereas, VAE ignores objects. Samples were generated at temperature of = 1 without cherry-picking.
These illustrations collectively showcase that flow vectors in the visual scene are shaped both by local variables, such the object's physical location, and their interactions with chobal-inftuencin variables like self-motion. By comprehending data across scales, cNVAE is able to capture such hierarchical and multi-scale properties inherent in naturalistic optic flow patterns. In contrast, hierarchical VAEs exhibit a complete failure this regard.
10.4 All models achieve good reconstruction performance
All models considered in this paper are trained based (in part or fully) on reconstruction loss, corresponding how well they could reproduce presented stimulus. Here report the models' reconstruction performance, to ensure that they are all able to achieve a reasonable performance. Figure [4 illustrates how the reconstruction loss (EPEPD, endpoint error per dim) and negative evidence lower bound (NELBO) depend on different values of 3 and model types. Both hierarchical and -in-hiirarchhcal VAE models perform well, with cNVAE exhibiting slightly superior performance for 3 < 0.8.
10.5 Untangling and reconstruction loss are strongly anti-correlated for eNVAE but not VAE
In this work, we demonstrated that virtue of hierarchical VAEs is that they make information about ground truth factors easily (linearly) accessible in their latent representations. We were able to quantify this relationship because used a simulation (ROFL) where ground truth factors were available. However, such a condition is not necessarily the case in real-world datasets that typically lack such direct knowledge 22
To address this, we examined the relationship between informativeness and reconstruction loss, as well as NELBO, depicted in Fig. [5] to test whether there is relationship between them. Strikingly. we observe a significant and strong negative linear relationship between reconstruction loss and informativeness in eNVAE (r -0.91, p 8.9 x 10-7 (-tts.). Conversely, such correlation is observed for VAE (r = 0.10, p = 0.72, test). Notably, this relationship is absent for both models
when comparing informativeness with NELBO (Fig. 15 These findings suggest that reconstruction
28

---

2g = -2.52
Zg m -1.85
Zz -1.17
-0.50
= 0.17
Zg 0.84
Zz 1.51
Figure 13:
Latent traversal performed using latent variable that effectively captures Youy- This
variable belongs latent group operating the spatial scale of 2 (see Figs. and table
loss exclusvvuyvv in the case of cNVAE can serve viable proxy for informativeness, even in the absence of ground truth factors real-world datasets.
10.6 Brain alignment requires more than just linear regression performance
In the field neuroscience, linear regression R² scores have commonly been employed measure
aligmment between biological and artificial neural networks [30 31 34 36). However, we found
29

---

10
EPEPD
10
model
cNVAE
VAE
10-3
cNAE
AE
10a
NELBO
x 10
4 10
x 101
ae 0.01 0.1
0.15
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
1.0
1.5
2.0
5.0
10.0
B
Figure 14: Both models demonstrate robust reconstruction performance, with cNVAE exhibiting slight improvement. EPEPD, endpoint error per dim (lower is better). NELBO, negative evidence lower bound (lower better).
cNVAE
cNVAE VAE
VAE
-0.91, = 8.9e-07
== -0.15, = 0.58
Informativenest
0.10, p = 0.72
-0.31, = 0.24
0.6
0.4
0.2
0.00
0.05
0.10
0.15
0.20
60
80
100
EPEPD
NELBO
Figure 15: Strong anti-corellation between reconstruction loss and informativeness (i.e., untangling) observed only for cNVAE. Each seatter point corresponds to model trained with different 3 value. EPEPD, endpoint error per dim (lower is better). NELBO, negative evidence lower bound (lower better).
that relying solely on regression performance provided weak signal for model selection, which is
consistent with recent work (36 112 125). In contrast, our alignment score defined in Fig. g
provided much stronger signal for differentiating between models (compare Fig. and Fig. 8c)
30

---

To investigate this further, we conducted statistical tests, comparing pairs of models based on either their performance in explaining the variance of MT neurons or their brain- -ainnmal scores. The results are presented in Table 5 We observed that ridge regression performance was not particularly effective in distinguishing between models. We measured the effect sizes between distributions of N = 141 neurons, measuring how well ridge regression performance (Fig. (7) or alignment scores (Fig. ( discriminated between models. found that the effect sizes were substantially larger for the alignment scores (Fig. [6). These findings suggest that while model performance is necessary, it not sufficient for accurately measuring alignment between artificial and biological representations. In conclusion, refined and more comprehensive approaches required to properly capture brain alignment 112 115
Test Performance Alignment
Effect size (Cohen's d)
1.0
0.5
0.0
005
0.01
0.1
0.15
0.3
0.5
0.6
0.7
0.8
0.9
1.0
1.5
2.0
5.0
10.0
Figure 16: Effect sizes shown for the two different approaches in model comparisons: using ridge regression Performance (related to Fig. 7]. versus, Alignment scores (related to Fig. [8)). Alignment demonstrates substantially larger effect size, leading more pronounced differentiation among the models. Positive values indicate larger value for hierarchical models (i.e., cNVAE > VAE cNAE AE). For different / values, model Performance difference fluctuates between positive and negative values. In contrast, the difference model Alignment scores is consistently positive, indicating greater brain alignment for hierarchical models.
10.7 Dissecting cross-validated model performance on neural data
One measure of neural alignment that we consider here used unsupervised latents to predict exper- imentally recorded neural responses (via linear regression). The neurons in this dataset had long presentations of optic flow stimuli, but only fraction of the neurons (105 out of 141 neurons) had multiple repeats of short stimulus sequence. Because these repeats facilitated much clearer mea- sure of model performance, due to the short stimulus sequence (~~5 sec; see Fig. (6)) used to estimate performance the neurons with stimulus repeats, here we perform additional tests to verify that our conclusions hold using more rigorous cross-van cross-validation procedure.
used these where available, following previous work (34). However,
Specifically, here we used the long stimulus sequence to both train model parameters (80% of data) and determine hyperparameters (ineluding regularization and optimal latencies) using a separate validation set (20%). We then predicted model performance on the repeats. This is in contrast to our approach in the main text, which used the repeat data itself to optimize hyperparameters, but as a result, risked overfitting on the short stimulus sequence. Indeed, we found that independent determination of hyperparameters (from the validation set) resulted in decreased model performance on the test set (Table for both cNVAE and VAE models, although with comparable effect sizes. As a result, our conclusions overlapping performance between cNVAR and VAE with slightly better performance of the cNVAE were robust.
31

---

Performance (Figure
Alignment scores (Figure 8c)
β
effect size
value
significant?
effect size
p D-vllue
significant?
-0.33
0.01
1.2
8e-29
0.01
-0.28
0.002
0.32
0.0004
0.1
-0.47
2e-07
/
0.12
0.2
X
0.15
-0.16
0.07
x
0.46
4e-07
/
0.2
-0.14
0.1
x
0.33
0.0003
V
0.3
-0.012
0.9
x
0.69
4e-13
0.4
-0.15
0.09
X
0.65
7e-12
/
0.5
-0.013
0.9
X
0.65
7e-12
/
0.6
0.29
0.001
0.69
4e-13
/
0.7
0.16
0.08
X
0.71
le-13
/
/
0.8
0.44
1e-06
0.73
4e-14
0.9
0.24
0.008
/
0.96
le-20
/
1.0
0.0016
X
0.89
1e-18
/
1.5
0.17
0.07
X
1.2
3e-28
2.0
0.037
0.7
X
0.91
3e-19
5.0
-0.24
0.008
/
1.3
7e-31
10.0
-0.056
0.6
X
0.55
3e-09
Table 5: Hierarchical architectures significantly more aligned with MT neurons compared to non-hiirarchical ones. Paired t-tests were performed, and p- valuus were adjusted for multiple comparisons using the Benjamini-Hocheee correction method [138). Effect sizes were estimated using Cohen's [139), where positive values indicate cNVAR > VAE (or cNAE AE, first row).
Performance, R (n se; N = 105)
Model
8 = 0.5
0.8
=
cNVAE
.452 + .025
.479 + .025
.486 = .024
.422 + .025
VAE
.461 + .027
.401 + .029
.466 = .028
.461 + .027
cNAE
.460 + .026
AE
.412 .031
Table 6: Model performance on the subset of neurons for which we had repeat data (N = 105). Here, we used repeat data as a held-out test set. We found that model performances remain statistically indistinguishable providing additional support for main results.
11 Additional Discussion
11.1 Visual-only cues and unsupervised learning are sufficient for untangling optic flow causes
It is well known that an observer's self-motion through an environment can be estimated using optic flow patterns alone [80]. although has been proposed that the brain also uses combination of visual and non-visual cues [40), such as efference copies of motor commands [41], vestibular inputs
32

---

[14], and proprioception [143). Indeed, using optic flow alone can be complicated by rotational eye movements and independently moving objects [48 145]. How can the brain estimate the motion of an object when self-motion is non-trivially affecting the entire flow field [16 147]. including
that of the object itself? This a critical motion processing problem [148), manifested in diverse ecological scenarios such as prey-predde prey-predator dynamics, crossing the street, or playing soccer. Are
retinal visual-only cues sufficient to extract information about different self-motion components, as well as object motion [
Unlike the approach of Mineault et al. 34 our work is based on the premise that retinal-only signals contain sufficient information about their ground truth causes and that these factors can be extracted using unsupervised learning alone. In contrast, Mineault et al. (34) trained their DorsalNet by providing dense supervision on self-motion components, arguing that such supervised training is biologically plausible since these signals are approximately available the observer from vestibular inputs and other such extra-retinal sources.
Our results (Fig. 4) provide proof-of-concept that it is possible to untangle ground truth causes of retinal optic flow in a completely unsupervised manner, using the raw optic flow patterns alone. Comparing the performance of alternative models revealed two essential ingredients for the success of our approach: (1) an appropriate loss function-inferent. U and, (2) a properly designed
architecture hierarhiccal, like the sensory cortex
11.2 Comparing eNVAE to NVAE, its non-compressed counterpart
In Fig. investigated the performance of the original non-compressed NVAE with approx- imately matched latent dimensionality (440, slightly larger than 420 for cNVAE: more details are provided in section 9..) We found that although NVAE did better than non-hierarchical VAE, it still underperformed cN INVEE. This is because trying to match the latent dimensionality of NVAE with cNVAE necessitates reducing the number of hierarchical latent groups in the NVAE since most of its convolutional latents capture spatial dimensions rather than more abstract features as cNVAE does. From previous work (7I). we know that the "stochastic depth" of hierarchical VAEs is the key reason behind their effectiveness, therefore, it was expected that reduced depth would hurt an NVAE with a matched number of latents.
is worth noting that eNVAE and NVAE had a balanced performance across all ground truth factors (Fig. D in that they captured both global (elfeem and local (object-related) aspects well. In contrast, non-hierarchical VAE completely ignored object-related variables and focused solely on the global scale determined by fixation point and self-motion. This suggests that hierarchical architecture is crucial for a more comprehensive understanding of multi-scale datasets such as ROFL.
11.3 Disentanglement is in the eyes of the beholder
The conventional definition of disentanglement in machine learning relies on relating the latent representation predefined generative factors, choice that is often not uniquely determined, and is often guided by heuristics (18). For example, in the present work, we chose to represent velocities in Cartesian coordinates, rather than polar coordinates and judged the degree of disentanglement based on the relationship of the latents to these factors. Notably, these coordinate systems are related by a nonlinear transform, and latents linearly related to one coordinate system will then be less linearly related to the other.
Throughout our experiments, we observed one more cNVAE latent variables that exhibited equally high mutual information with all components of self-motion velocity in Cartesian coordinates (c.g.. latent dimensions 154 and 122 in Fig. 3) When we translated the generative factors into polar coordinates and looked at their relationship with these latents, we found that these latents were solely encoding the magnitude of self-motion, irrespective of its direction (Fig. L7p. Thus, these seemingly highly entangled latent variables would have achieved almost perfect disentanglement scores had we employed polar coordinates instead of Cartesian coordinates to represent the ground truth velocities (d, s 0.38 + d x 0.81).
Our findings shed light on the fact that disentanglement cannot be considered an inherent property of objective reality; rather, is significantly influenced by our a priori determination of independent factors of variation within the data. Given the inherent complexity and richness of natural datasets,
33

---

d 0.38
0.81
1.0
M
0.5
0.0
Fy
V SeeFyy
self,
Xobi
obj
Y obj
V
Lobi
^
doll
x
dep
obj
Vobj.x
2 /
V
X
Figure 17: This latent variable cares about the magnitude of self-motion only. Its disentanglement score depends strongly what choose priori our independent factors of variation in data.
becomes challenging to unequivocally determine the true factors of variation priori, pointing the need for more systematic approach.
When it comes to interpreting brain functions, perhaps considering the behavioral relevance of generative factors could guide our choices. For example, an animal directing its own movements does not pick grid coordinates but rather chooses direction and speed, possibly making polar representation more "behaviorally relevant" to representations within its brain.
Another promising approach lies in focusing on the symmetry transformations inherent in the world [18 [50]. By seeking representations that respect the underlying geometrical and group-theoretical structures of the world [15]] [52]. and comparing them in principled ways (112 114). we may uncover new perspectives that provide valuable insights beyond the conventional and subjective notions of disentanglement.
12 Additional Background & Related Work
The deep connections between VAEs and neuroscience is explored in review by Marino [17]. work looking at the synergy between neuroscience and VAEs has thus far been limited to the following recent papers.
Storrs et al. [15) trained PixelVAE 153 models synthesized images of glossy surfaces and found that the model spontancously learned to cluster images according to underlying physical factors like material reflectance and illumination, despite receiving no explicit supervision. Most notably, PixeIVAE mimicked human perceptual judgment of gloss in that it made errors that were similar to those made by humans.
Higgins et al. [16] trained 3-VAE 85 models on face images and evaluated the models on predicting responses of individual inferotemporal (IT) neurons from Macaque face patch. They found that 3-VAE discovers individual latents that are aligned with IT neurons and that this alignment linearly increases with increased beta values.
Csikor et al. (68) investigated the properties of representations and inference a biologically inspired hicrarchical VAE called Topdown VAE that is capable of reproducing some key properties of representations emerging VI and V2 of the visual cortex. The authors showed that architectures that feature top-down influences in their recognition model give rise to richer representation, such that specific information that not present in mean activations becomes linearly decodable from posterior correlations.
Finally, Miliotou al. 669 learned mapping from the latent space a hierarchical VAE to fMRI voxel space. This allowed the use of the trained decoder to reconstruct images directly from brain responses. Importantly, they perform ablation experiments and find hierarchy an essential component of their models.
34

---

13 Appendix: ROFL derivations and detailed description
ROFL aims generate synthetic optical flow patterns that closely resemble those arising from ecologically relevant types of motion. This entails the consideration of both self-motion and object- motion components. Specifically, the self-motion component encompasses various subtypes such as translational and rotational motions induced by eye or head movements. Furthermore, it is crucial to have control over the relative prominence of these two types of motion. Depending on the specific object- deneiated characteristic when an
configuration, the resulting flow pattern exhibit rel-mootoo--dominteed nature if fewer or no abundance objects is incorporated, or conversely, a self-motion objects present. By carefully manipulating these factors, ROFL enables the generation of optical flow patterns that faithfully capture the dynamics of real-world scenarios involving an interplay between self-motion and object motion.
13.1 Setup
Self-motion is meaningful if the observer is moving relative to a static background (i.c., the environ- ment). We consider a solid background wall distance from the observer which extends to infinity in the X Y plane. Define the following coordinate frames
1. Fixed coordinates, represented with capital letters: (X,Y
2. Observer's coordinates, represented with lowercase letters: (æz, y, 2).
Both coordinate systems are centered around the observer. The difference is that is always perpendicular to the background plane, while 2 is defined such that it is parallel to the fixation point.
All points in the space are initially expressed in the fixed coordinates. For example, any given point like P (X, Y, Z) has coordinates defined in the fixed system. The same is true for fixation point, F (Xo, Yo. Z), and the observer's position Ô = (0, 0, 0), and so on. See Fig. 18
13.1.1 Relating the two coordinate systems
How do we go from the fixed coordinates to observer's coordinates? We need to find a rotation matrix that maps the direction to the direction defined by the fixation point. By definition, this will determine å and all the other directions in the observer's coordinates. First, let us express F using its polar coordinates:
F = (Xo, Yo. b) = -11211, Oo 2o).
(19)
where ||
22, cos Ou = F|| and tan P = Vo/Xo, It turns out what we are looking for rotation by an amount O₂ about unit vector that depends only on Py. That is,
X
sin Po
a($b)
cos to 0
(20)
and the rotation matrix is given by Rg (On) + (sin θa) L(ū) + cos θa) L())2.
Here, L(W) is the skew-symmetric matrix representation of a vector W, that is
L(ū) = 3 -- wx Lx + wy + wz wzLz, where L; are the familiar basis matrices for s3)),, the Lie algebra of 50(3):
(21)
0 0 0) - Lx =
0
-1 0
Ly =
0
Lz =
0 0
(22)
1
0
0
0
The rotation matrix in its full form is expressed
35

---

a
Fixed coordinate system:
(X, yz2
Observer's coordinate system:
F
2
F (Xo; Yo, 1)
Observer
Fixation point:
00
Z
P
Any given point in the scene:
p= (X, Z2
Solid background
Y
Ô Ō (0,0,)
y
Observer:
These points are all defined in the fixed coordinates
b
C
P
Z
F
X
F
θ
O
Go
40
Z
Observer
X
F
Oy
Y
Y
U
tan Otr = tan COS Ø
tan City tan sin Ø
Figure 18: Setup. (a) Background and coordinate systems. (b) The rotation of the fixed coordinates
(x, 8,2 2) onto the observer-centric coordinates (á, 0. 2) can be accomplished by applying rotation of angle O around the unit vector as defined Eqqattinn 20 Note that by definition, é always parallel to the fixation point. See the text for more details. (c) Perspective from the observer's point of view. The retinotopic coordinates of point P are denoted by a, and Cty, measured in radians.
82 + c² cos Oo -cs cos Oul cssn Ou
1 cos Op]
esin Ou sin Og cos O
R= Rg(
c² + s² cos O
(23)
ssnn Ou
where we have defined c = cos o and = sin o
Any point P (X,Y,Z) defined the fixed coordinate system can be mapped to point p R P with components ş = (xz, y, 2) given in the observer's 's system. For example, it can be shown that Rd F is parallel to 2, which is how we defined R in the first place. See Fig. 18
In conclusion, we have rotation matrix R that fully determined function of the fixation point F. Applying R rotates the fixed coordinate system onto the observer's system
a = RX, ÿ RÝ, 1 == R².
(24)
This way, the observer's coordinates (z. y.2) 2) are obtained ensuring that the fixation point always parallel to &.
13.1.2 Adding motion
Consider a point F = (X, ,, 2). In the previous section, we showed that the rotation matrix in
Equation 23 can be used to relate P to its observer-centric coordinates
36

---

p= Rp.
(25)
Differentiate both sides to get
dP
R dp di
(26)
di
dt
where V the velocity of the point relative to the observer, with components measured in the fixed coordinate system.
13.1.3 Derivative of the rotation matrix
Here we will derive closed-form expression for dR/dt in terms of the observer's 's angular velocity. This is important because it essentially captures the effects of eye movement.
Let us use a (wx,, wy, wz) to denote the angular velocity of observer's coordinates, with compo- nents expressed in fixed coordinates. The time derivative of R is given as follows
dR L
(27)
dt
where L(w) is given in Equation 21
Taken together, equations 3, 4, 8, and 9 yield
dp RTV
[R ((J)
di
0
Wy
(28)
RTV
Rd
W2
W
R P
Wy
wX
This equation allows relating the velocity of moving point expressed in fixed coordinates, V. to its observer-centric equivalent, T = dø/dt. dt. Importantly, Equation 28 is a general relationship valid for any combination of velocities and rotations.
Eqquttion
13.1.4 Motion perceived by the observer
Before discussing specific applications, we need to solve one remaining piece of the puzzle: what is J. and how is it related the motion perceived by the observer? To address this, need to map the retinotopic coordinates the observer (measured in radians), to points in the environment (measured in the observer-centric coordinates).
Denote the retinotopic coordinates using (Or,, aty). The goal is to relate each (ar, (gy) pair to a point in the world with coordinates (xr, y, 2). It can be shown that
lan 0g x// lan O cos Φ
(29)
tan 0g = y/z -tan sin Φ, where 8 and & are the polar coordinates p. See Fig. 18c for an illustration.
Differentiate both sides of Equution 29 and rearrange terms to obtain the time derivative (azz, agy) function of ह (ir, ú, ø) and p. That is,
z2 + x
(30)
yz
ży
dy
z² + y2
Thus, we obtained a relationship between retinotopic velocity (cin;(y)) (measured in radians/ s) and velocity of objects in the world (measured m/s). In conclusion, equations 28 and 30 enable us to compute the observer's perceived retinotopic velocity for any type of motion.
37

---

13.2 Example: maintaining fixation
We are now ready to apply the formalism developed the preceding section create optical flow patterns in realistic settings. In this section, we will consider specific example: self-motion while maintaining fixation on a background point. This an important and ecologically relevant type of motion, as we often tend to maintain fixation while navigating the environment- -thnk about this next time you are moving around. In fact, one can argue that fixation occurs more naturally during behaviors such as locomotion, as opposed to typical vision experiments where the animals are stationary and passively viewing a screen.
What does it mean to maintain fixation while moving? It means the eyes rotate to cancel out the velocity of the fixation point. Assume the observer is moving with velocity Veety. Equivalently, we can imagine the observer is stationary while the world is moving with velocity V =-Veet (this is the same V that appears Equation 28
Any eye movement can be represented mathematically as rotation of the observer's s coordinates with respect to the fixed coordinates. Use to denote the angular velocity of the rotation. We can determine 3 by requiring that the rotation cancels out the velocity of the fixation point. To this aim, first compute the normal component of V with respect to the fixation point F as follows
F
(31)
11
Use the above equation compute the angular velocity
Fx
V
W
(32)
FF12
28to Plug Equation 32 in Equation 28 to find o7, and use Equation 30 to find (áz, dy). This way, we just calculated the motion perceived by the observer during self-motion while maintaining fixation.
13.3
The algorithm
Start by choosing a fixation point F = (Xo, Yo. 2), field-of-view extent (fov), and retinal resolution (res). For example, for = 45 and res = 5.625.. Use this to create mesh grid for (axr, a.y). The grid covers (-for, ffo) degrees in each dimension and has shape 2 fov pov/res + 1, or 17 x 17 in our example. Use U/seEEquation 29 to compute (6, Φ) pair for each point on the (0z:0y) grid.
Because the simulation is icale-invariant, we can always fix the distance between the observer and the background wall. Assume Z = de. 1, and equations 23 25 find ş (x,, y 2) per grid point. Note that we assume infinite resolution for the points in the environment.
So far, partitioned the retina into grid and found the observer-centric coordinates of point in sample a random self-motion velocity Vaey and use the results described in the previous sections to find (di,, dy). This provides us with the 2D retinal velocity vector the grid and concludes the algorithm for the simple case of fixate-0 (i.e., no objects).
the real world that falls each grid point. Now
Please refer to our code for additional details, including instructions on how to incorporate indepen- dently moving objects into the scene.
13.3.1 Valid fixation points
The observer-wald distance is constrained to be positive value. That is, the observer is always situated on the left side the wall (Z > 0 for all points; Fig. I8). This constraint, combined with a given for value and and Eqquation 23 results in theoretical bound which fixation points are valid. To derive this bound, we start by considering the Ž component of Equation 25
(Rp) RaI Razy + Rag2 Z 0.
(33)
38

---

Divide both sides by and use Equation 29 to get
Ryi tan o + Ra2 tan / + Ra3 >> 0.
(34)
Plug the matrix entries from Equation 23 and rearrange get
sin Op(cos Py tan a + sin t tan () < Og.
(35)
Xo tan & + Yo tan B < Z.
For a given fov value, we have - tan (fov) < tan a, tan 8 < tan (fov). Assuming 2 = 1, we find that the (Xo, Yo) coordinates of valid fixation point should satisfy the following inequality
Xul
[Yl (fov))
(36)
In other words, the L, norm of (Xo, Yo) should be less than 1/ tan (foc). For example, when foo = 45", the upper bound 1. This result can used to choose an appropriate for: smaller values allow a larger phase space to explore for fixation points.
13.4 Conclusions
We introduced Retinal Optic Flow Learning (ROFL), a novel synthetic data framework simulate realistic optical flow patterns. The main result is given in Equation 28 which shows the simulation has the potential generalize engage in smooth pursuit of a moving object instead of maintaining fixation, a scenario that can be modeled simply by adjusting the choice of w in Enqution 28 This flexibility enables us to generate diverse repertoire of realistic optical flow patterns, which can be used in training and evaluating unsupervised models, among other applications.
various other types of motion. For instance, the observer might
We believe that our framework will prove valuable not only to neuroscientists studying motion perception but also machine learning researchers interested in disentangled representation learn- ing. By bridging these domains, we aim to foster interdisciplinary collaboration and advance our understanding of both biological and artificial information processing systems.
39


Reviews:
[{"review_id": "3qq5GSHxbK", "reviewer": "Reviewer_Hbpj", "summary": "The paper investigates the alignment of representations in deep generative models with activity in mammalian nervous systems. They provide a novel dataset on motion perception (Retinal Optic Flow Learning or \"ROFL\") against which to test computational models of Helmholtzian analysis-by-synthesis. The dataset generates retinal flows based on disentangled latent factors determining the motion and appearance of objects.  The paper then tests a novel extension of the deeply hierarchical Nouveau VAE which imposes a \"pyramidal\" scaling of latent spaces through the hierarchy for its ability to capture the retinal-flow dataset in a \"Helmholtzian\" way, as well as its alignment with a pre-published macaque electrophysiology dataset on the middle-temporal visual area (MT).", "questions": "Do the authors have a plan to apply their cNVAE in a \"clockwork VAE\" setting to predict videos?  Can we see if optic flow emerges from training a deep generative model?  Moreover, is there an ablation study able to show how the recognition model $q$ versus the generative model $p$ in the cNVAE contribute to its alignment with neural representations?\n\nThe authors have presented interesting ablation studies in their rebuttal, whose results make perfect sense in retrospect and yet which I could not have predicted a priori.  I am thus convinced I need to raise my score.", "limitations": "The authors have addressed potential limitations and impacts.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The paper presents an original deep generative model architecture, trainable with the typical ELBO or $\\beta$-VAE objective by black-box variational inference, which captures a desirable and sought-after feature of generative modeling in the brain.  The experimental evaluation compares against maximum-likelihood/reconstruction-loss estimation and a deterministic autoencoder, a nice comparison for making sure that the neural data aligns best with a deep probabilistic generative model as opposed to a non-probabilistic deep neural network.", "weaknesses": "As the authors admit in the Discussion section, they trained their compressed Nouveau VAE (cNVAE) on optical flow data rather than on video/pixel data.  While this does map better to the MT area than pixels do, a very deep model ought to be able to learn to represent optic flow in the course of predicting pixels.\n\nWhile this supervision with feature-engineered optical flow data in place of learning from time-series does remain a weakness, the authors have presented a sufficiently interesting additional contribution in addressing my question below that I will be raising my score.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> As the authors admit in the Discussion section, they trained their compressed Nouveau VAE (cNVAE) on optical flow data rather than on video/pixel data.\n\nWe appreciate this point, and it relates fundamentally to key choices in our approach. The raw visual image (photons on the retina) is processed by many layers of the visual pathway through many specialized computations throughout the retina and cortex, and these transformations are the subject of much computational work (e.g., [Nishimoto & Gallant](https://www.jneurosci.org/content/31/41/14551)). Here, we wanted to simplify the stimuli to focus on a clearly defined and focused question of optic flow: starting with the optic flow (velocity field resulting from movement in the scene and by the observer/eyes) to test if our artificial model can extract the generative “ground truth” variables from the optic flow stimulus. This stimulus also purposefully matches a majority of experiments in area MT that use moving dots, including the dataset in our study.\n\nThe use of optic flow instead of raw-pixel stimuli of course allows the cNVAE to focus on the problem of inferring a representation of the generated motion fields: and, in this sense, not need to simultaneously solve the difficult problem associated with extracting motion fields from pixel data (and/or the conjoint problem). In short, our manuscript focuses on this problem – and please see our summary of the general goals of this manuscript – and, in achieving high performance on explaining MT data, implies that MT may be doing something similar (when given equivalent dot-motion stimuli).\n\nIn terms of solving the harder problem of processing raw spatiotemporal movies, we propose our approach provides a new possible path for this and plan to extend our cNVAE to operate on raw images in the future.\n\n> While this does map better to the MT area than pixels do, a very deep model ought to be able to learn to represent optic flow in the course of predicting pixels.\n\nYes, in the future we will test this. We expect such solutions based on raw pixel data will be achievable with the same framework and more layers of latent variables, although it is possible that these problems are linked.\n\n> Do the authors have a plan to apply their cNVAE in a \"clockwork VAE\" setting to predict videos? Can we see if optic flow emerges from training a deep generative model?\n\nThank you for this question. In the Clockwork VAE paper, they use hierarchical latents but in the temporal domain and report the benefits of this inductive bias. In our study, we decoupled spatial and temporal aspects of motion processing and focused solely on the spatial hierarchy. This simplification allowed us to demonstrate that hierarchical models were able to understand multiscale data (like the real world) where a non-hierarchical VAE struggled.\n\nIt would be interesting to add temporal hierarchical structure to ROFL in a spatiotemporal video setting and use it to investigate how latents would capture coexisting fast and slow dynamics. For instance, one can include objects that move much faster or slower than self-motion. In the clockwork VAE paper, they found that higher latents captured objects with slower time scales. Here, we found that higher latents specialize in encoding object-related features. It would be interesting to combine their setting with our architecture to find out what happens in the spatiotemporal domain and whether the model learns brain-like representations. This is a scientifically relevant and important question because there is also a hierarchy of intrinsic time scales in cortical dynamics ([Murray et al., 2014](https://www.nature.com/articles/nn.3862)). However, we believe such an important problem deserves its own full treatment in a separate, follow-up paper.\n\n> Moreover, is there an ablation study able to show how the recognition model q versus the generative model p in the cNVAE contribute to its alignment with neural representations?\n\nThank you for this excellent suggestion. We performed ablation experiments and found complementary insights into why cNVAE outperforms alternative models. We report these results in Figure R1, along with a brief discussion and interpretations in our global rebuttal. Overall, we found your suggestion to be very interesting because understanding the contributions from bottom-up and top-down connections in the cortex is a central problem in neuroscience. We are planning to tackle this question more systematically in the future.\n\n> The authors do not appear to include a Broader Impacts section to review.\n\nWe had to fill it out here on OpenReview rather than having a designated section in the paper. Our response is copied below for the reviewer’s consideration:\n\n\"**Broader Impacts**: We introduce a new simulation framework that facilitates hypothesis generation and testing in science, but it is not sophisticated enough for potentially harmful applications. Thus we do not anticipate negative social impacts from this work.\""}, {"review_id": "ZqlqyFhynJ", "reviewer": "Reviewer_Y46x", "summary": "The authors present a framework to evaluate motion detection in different DNNs. First, they intoduce a new concept to create flowfields for optical sitmuli, which include local and global motion and additionally fixation points. They use the parametrized stimuli to train a new hierarchical VAE (cNVAE) and compare it to other DNNs in terms of different disentanglement metrics and similarity to neural recordings of MT. cNVAEs outperform standared VAEs in the disentanglement setting, and are comparable to VAEs for neuronal prediction while their latent is more aligned with single neural recordings. Training on different synthetic flowfield datasets showcase the influence of the treaining data to predict neural respones. ", "questions": "- Which $\\beta$ is used for Fig. 3?\n- Can the authors reiterate how cNVAE is used for data generation? \n- How does cNVAE compare to the non-compressed version (with matched latent dimensions)?\n- How do the different metrics perform on the raw data? This would add a simple baseline which could help to interpret the results.\n- Fig 4 a: from which layer do the neurons come? Maybe even highlight them in Fig. 3.\n\n", "limitations": "Yes, very well done.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Clear description of stimulus genearation, and overall a solid paper. \n\nClear and detailed description of the limitations. \n\nInteresting hierarchical VAE model which could be transferred to other learning tasks and brain areas. ", "weaknesses": "**Major**\n- Evaluation: While the authors do a good job in a high-level evaluation of the models (applying different evaluation scores and comparison to other models), further in-depth investigation would strengthen the paper significantly.\nFor example, it would be interesting to investigate: the learned representation (manifold); the necessary dimension of the latent space, and how the performance and disentanglement change with different dimensions; the hierarchical structure (in the model and the neural data); counterexamples where the model fails and why it fails; further investigation of the receptive fields for different neurons; some of them are touched on in the Appendix, but a clear link to the main text is missing. \n- The paper would benefit from some restructuring to disentangle previous work and Methods. Each subsection in Section 2 seems to start with a short introduction and previous work.\n- l. 108 ff. The model architecture (especially the processing of the sampling layers) should be described more clearly. \n- It is not clear why the authors chose a specific $\\beta$ in different places. (For example in Fig. 4; Fig. 8b seems to be cherry-picked for cNVAE, and an especially bad example for VAE)\n- All Figures: more descriptive captions would be beneficial.\n\n**Minor**\n- The authors could reference to specific sections/figures in the appendix to make navigation easier.\n- The layout of the tables does not follow the style guidelines for scientific tables (see NeurIPS guidelines)\n- Fig. 4a is not mentioned or described in the text. \n- Fig. 4b: mention R^2 score in caption. Write R^2 numbers down also for “bad” models. It is not clear why some bars are “missing”.\n- The authors should include a short description of the DCI framework.\n- Fig. 6: figure labels are very tiny and hard to read, 6d is so tiny that it is hard to interpret it at all. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Evaluation...\n\nWe like these suggestions and find some particularly exciting. Although, we consider several of them interesting future directions, given the manuscript already covers a lot of ground (see general response for a summary). In short, the present work is meant as an empirical report of all the cool things that happen to the representations when the latent space is designed in a hierarchical way. As a next step, we were planning to quantify the geometry of the latent spaces using tools from neural population geometry approaches that have recently gained traction.\n\n> the learned representation (manifold)\n\nMotivated by your comment, we performed some first-level analyses to investigate the geometry of representations using the simple method of “effective dimensionality”, which is computed based on eigenvalues of the covariance matrix and provides an estimate of manifold dimensionality. We found that across a broad range of $\\beta$ values, the dimensionality of cNVAE representations was substantially larger than that of VAE, suggesting that their representational geometries are ultimately different in a quantifiable way. This result be taken as a starting point for a thorough analysis of the representational geometries. Please see our Figure R1 and global rebuttal for more details.\n\n> the necessary dimension of the latent space\n\nIt is interesting to quantify the dependence of our results on the number and organization of latents (i.e., how many latent groups, how many latent variables per group, etc). We did not perform these analyses for the rebuttal due to lack of time, but if the reviewer feels like this would enhance the quality of our paper we will do so and include the results in the final version. Please let us know!\n\n> counterexamples where the model fails and why it fails; further investigation of the receptive fields for different neurons\n\nWe will also explore more neurons and find counterexamples to report in the supplementary. We will dig a bit deeper to understand why some neurons are more aligned, while others are not.\n\n> The paper would benefit from some restructuring\n\nThank you for the suggestion, which we agree with. We will add a “related work” section and clarify some of the missing key background info, which is added to our global rebuttal.\n\n> l. 108 ff. The model architecture (especially the processing of the sampling layers) should be described more clearly.\n\nWe apologize for the lack of clarity and will include more details in a substantially revised description of the model architecture and our specific contributions.\n\n> It is not clear why the authors chose a specific beta in different places. (For example in Fig. 4; Fig. 8b seems to be cherry-picked for cNVAE, and an especially bad example for VAE)\n\nOverall, we (and others) find that different choices of the $\\beta$ will result in different model properties. As a result, while we scanned across all betas, we displayed results for certain betas to best demonstrate our points. For instance, in Figures 3 and 4 we chose beta values that maximized the overall informativeness score for each architecture in order to make a fair comparison ($\\beta = 0.15$ for the cNVAE, $\\beta = 1.5$ for the VAE -- please ignore the typo in Figure 4 caption, it should say $\\beta=1.5$, we will fix this). You can see that this is the case from Figure 5 where we show DCI scores for all betas. In Figure 8b we deliberately chose a larger beta for VAE because previous work ([Higgins et al. 2021](https://www.nature.com/articles/s41467-021-26751-5)) suggested that increasing $\\beta$ values increases alignment, which we did not observe here. In contrast, even a very small $\\beta = 0.01$ for cNVAE results in a large alignment. This result (paired with other observations) suggests that 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$ values, so ultimately a combination of both architecture and loss is important (but mostly architecture).\n\n> All Figures: more descriptive captions would be beneficial.\n\nThis will be fixed in our revised manuscript, including the other minor suggestions (which we thank the reviewer for).\n\n> Which beta is used for Fig. 3?\n\n$\\beta = 0.15$ for cNVAE, and $\\beta = 1.5$ for VAE.\n\n> Can the authors reiterate how cNVAE is used for data generation?\n\ncNVAE was not used for data generation.\n\n> How does cNVAE compare to the non-compressed version (with matched latent dimensions)?\n\nWith a matched number of latents, it is expected that NVAE will severely underperform cNVAE. This is partly because trying to match their latent dimensionality necessitates reducing the number of hierarchical latent groups in the NVAE. From previous work ([Child 2021](https://openreview.net/forum?id=RLRXCV6DbEJ)), we know that the “stochastic depth” of hierarchical VAEs is the key reason behind their effectiveness; therefore, we expected that reduced depth was going to hurt an NVAE with matched # of latents.\n\nTo test this, we trained a non-compressed NVAE with roughly the same # of latents (440 vs. 420 for cNVAE), the same number of parameters and conv layers, but necessarily with a reduced number of latent groups (11 vs. 21 for cNVAE). We tested its *untangling* performance (similar to Figure 4b) and found that it dropped significantly compared to cNVAE in predicting every ground truth variable, but it was still higher than VAE. The average untangling scores are as follows:\n\n- cNVAE: **0.898**\n- NVAE: 0.639\n- VAE: 0.406\n\nWe thank the reviewer for this question because it prompted us to demonstrate clearly that the latent space of NVAE is unnecessarily and redundantly large. An observation made by others as well ([Hazami et al., 2022](https://arxiv.org/abs/2203.13751)). We are planning to include this new result in Figure 4b, along with performance on raw data (as was also suggested) to further enhance the comparisons.\n"}, {"review_id": "Qd6glSe4qq", "reviewer": "Reviewer_ZQpT", "summary": "The paper introduces a synthetic data framework called Retinal Optic Flow Learning (ROFL) and uses that framework to test the performance of unsupervised models on two learning tasks: reconstructing grand truth variables and predicting the response of MT neurons. By imposing a latent hierarchical structure, the authors observed improvements along three axes: on the linear decodability of the ground truth, on the predictability of MT neuron responses and, finally, on identifying the causal structure of the world as a major factor driving these results.", "questions": "“Importantly, this framework allows us to manipulate both the architectures of the models and the causal structure of the world they are trained on” -> As far as I understood, you chose the models and the method for generating synthetic data. Being able to manipulate the architecture and structure of data is somehow expected. Why is that then an important fact and statement to make? If it’s not, please consider removing it. \n\n“We found that a single inductive bias, hierarchical latent structure, yields several improvements.” -> What is the metric that you use to assess performance and what do you compare against? The metrics are only qualitatively described in the paper. Please consider condensing that information in a clearly defined section, something like “Metrics and Baselines”, and give a clear, quantitative definition of the metrics. \n\n“the brain engages in hierarchical Bayesian inference ” -> Based on what I read, I do not believe you have presented sufficient evidence to back up such a strong claim. The experimental study on predicting MT response does not suffice. If you meant to back the claim using more of the results you presented please make it more clear in the main manuscript or else remove it. \n\nThe metrics that are used in the experimental section such as informativeness, disentanglement and completeness are not clearly defined. There are some papers cited where those metrics are mentioned in the text and I am sure that someone would find the definition if they look them up. However, given that these metrics are not widely used, I think the authors should give the definition in the text. \n\nWhat are the details of the data that you used for the MT learning task? There’s little to no information regarding what these data are, how they are collected, pre-processed etc. I understand that you cite a couple of papers that probably have these details, however, I think you have to give some details and explain how you use this dataset. \n\nWhat is the learning task for the MT prediction, i.e., what is the quantity you’re predicting? Is it the firing rate shown on figure 6? If so, how is that related to the hierarchical modeling in the previous section? Details on the dataset would help clarify. \n\nWhat is the AirSim dataset that the authors compare against in the MT prediction task? And how does it related to the newly introduced ROLF? If it’s a completely different synthetic data framework I do not see the point of comparing against it. \n \nIs the label in Figure 3 correct? There’s a reference to Figure 2. Maybe the author meant to refer to the bottom of Figure 3.\n\nOn lines 225-227 you say “all metrics for a broad range of β 227 values (Figure 5).” I am assuming that by all metrics you mean the  “untangling”, “disentengling”, and “brain-alignment” mentioned in line 195. However, the labels on figure 5 are slightly different, “Informativeness”, “Disentanglement” and “Completeness”. Please clarify what the metrics are and quantitatively define them, as they are not well known metrics.\n\nWhy the word “learning” in ROFL? There is no learning involved in the data generation, as far as I understood.", "limitations": "The authors have clearly and openly discussed the limitations of their work.", "rating": 7, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "As far as I understood (given the lack of clarity on a few aspects regarding the problem definition and related work, please look at my comment below), the authors  introduced a method to generate synthetic data and then modeled the generated data using an ensemble of model architectures of their choice. While I am dubious of the novelty and the sensibility of the idea, the experimental results presented seem to validate the argument. Another important strength of the paper is that the authors are pretty open and clear about the weaknesses of their method, which is something to comment them for.", "weaknesses": "I am dubious of the idea because it seems that the authors generate synthetic data and then choose an ensemble of architectures to model them. Also, it is hard to understand what is the main problem tackled by the paper, how this work relates to the state of the art and most importantly what is the improvement and contribution compared to the SOA. \n\nThere’s no Related Work section in the manuscript and the method is not compared against any other methods in the experiments. Are there many papers that deal with the same problem? Reference 29 seems to be one such paper and it is compared against in the experiments. However, given that this reference has a different dataset and model, I see no common ground for comparison. And given the peculiar structure of this problem area (generating data and then modeling them), the authors need to find way to highlight their contributions.\n\nAll these make it hard to assess the paper. I recommend the authors add a related work section, clearly connecting their paper with the state of the art, and clarifying the novelty and contribution. \n\nIn you related work section, you may want to consider referencing some papers that use hierarchical latent spaces for real-world data, such the following one (which is a newly published paper, so it’s understandable that you hadn’t included it).\n\nGenerative Decoding of Visual Stimuli, https://openreview.net/pdf?id=57OuafQmu8", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> As far as I understood…comment them for.\n\nWe struggled with how to answer this review. It seems that comments alternate between not understanding or knowing the relevant literature, and asserting (confidently) that the work was not novel and/or incorrect. Furthermore, in several cases, we already have whole sections devoted to some of the reviewer's concerns.   \n\nAs a result, we hope that we can largely address these concerns by clarifying our contributions, and have reworked parts of the text accordingly. Along these lines, we have included a draft Related Work section in the Global Rebuttal, and this should clarify where this paper sits in the existing literature. We have also reiterated the main points in our global response to reviewers and addressed individual ones here. We hope this clarifies the “sensibility” of the ideas, which stem from two major conjectures in neuroscience.\n\n> I am dubious of the idea…\n\nWe were motivated by the idea that hierarchical inference is important for learned representations. This is an old idea in neuroscience and a newer idea in ML. To evaluate the role of hierarchical inference in learned representations, it was essential to generate synthetic data with ground truth variables. We then evaluated the representations using a number of metrics and concluded that hierarchical inference leads to several improvements using metrics that have been proposed by neuroscientists and ML researchers.\n\nIn addition to our main goals, our paper moves beyond solely evaluating the reconstruction performance of hierarchical VAEs. As far as we know, a comprehensive investigation of representation learning in hierarchical VAEs has not yet been done: this is the first.\n\n> …and the method is not compared against any other methods in the experiments. Are there many papers that deal with the same problem?\n\nWe compare directly to a recently established benchmark, which we cite extensively: ref [29] – which was a spotlight NeurIPS paper in 2021. We are the only other paper to use this benchmark thus far.\n\n> Reference 29 seems to be one such paper…no common ground for comparison.\n\nWe use the same MT dataset as Ref 29 (crcns-mt1), and the model is indeed different which we now detail in the Related Work.\n\n> And given the peculiar structure of this problem area …\n\nSynthetic data with ground truth factors is essential to interrogate the learned representations and is standard practice in disentangled representation learning: https://paperswithcode.com/task/disentanglement\n\nOur manuscript had an entire section on this point (specifically, section 2.1 “Using synthetic data to test hypotheses about the causal structure of the world”)\n\n> In your related work section, you may want to consider referencing...\n\nThank you for this ref. We now address this (and other) papers in the Related Work section. We note here that Miliotou et al. have a very different goal than ours: to decode images from fMRI voxel activations (using a hierarchical VAE); whereas, our work focuses on learned representations and how they depend on architecture (hierarchical vs. non-hierarchical), loss function (variational vs. maximum likelihood), and the generative factors of variation in the training set.\n\n> Q1: As far as I understood...\n\nBriefly, our ability to manipulate the structure of the data (and model) is exactly why it is useful for our work. We highlight this point at multiple places in the manuscript, including in the Introduction and discussion (line 290). We even include a section about it (section 2.1 “Using synthetic data to test hypotheses about the causal structure of the world”). \n\n> Q2: What is the metric that you use to assess performance…\n\nBecause our focus was on the learned representations, the definition of the performance metrics based on the reconstruction loss was relegated to the Appendix (Fig. S5 shows a summary). We already did state in the Introduction (line 54) what we meant by “improvements”. We will add a table summarizing reconstruction performance and additional text describing the metrics computed.\n\n> Q3: “the brain engages in hierarchical Bayesian inference \"...\n\nHere we agree. Our observations are consistent with this notion, rather than prove it. We will modify the language accordingly.\n\n> Q4: The metrics that are used in the experimental section...\n\nThese metrics are widely used in the VAE literature, but perhaps not more generally. We will add more details about these metrics in the revised text.\n\n> Q5: What are the details of the data...\n\nWe will add some key details to the main paper, and refer the reader to the supplemental where we include every detail to make this work completely explained given the length constraints of the main text.\n\n> Q6: What is the learning task for the MT prediction...\n\nWe follow standard neural modeling approaches and predict the binned spike count of each neuron, which when normalized by the time bin size, results in a “firing rate” (number of spikes per second).\n\n> Q7: What is the AirSim dataset…\n\nBoth ref 29 and our manuscript analyze MT data from the same CRCNS dataset. We both train models using completely separate synthetic datasets that we then use to predict MT. Ref 29 used AirSim to train supervised 3D ResNets that they then use to predict MT (using a completely different stimulus). We created ROFL to simplify the causal structure of motion in the world and then, as a part of our analysis of the learned representations, we predict MT. Importantly, we get a factor of 2x performance gain over the previous best model (ref 29). \n\n> Q8: Is the label in Figure 3 correct?\n\nYes.\n\n> Q9: On lines 225-227 you say...\n\n“Untangling” and “informativeness” are the same which we discuss in lines 197-202 and also in supplemental section 1.4.\n\n> Q10: Why the word “learning” in ROFL? \n\nThe dataset is for unsupervised learning of latents: thus we consider this to be a “learning framework” overall.\n"}, {"review_id": "kcpYSeWKlD", "reviewer": "Reviewer_nG5y", "summary": "- This empirical study develops a new framework for neural representation of motion in cortex.\n- The authors propose a novel stimulus synthesis method for model training, parametrically generating optic flow fields w/ low-dim latent structure, rather than using pixel-space (i.e. image-computable) inputs.\n- They modify previously-developed hierarchical Variational Autoencoders, and argue that the inductive architectural biases of such networks produce a more desirable representation of motion, as quantified via entanglement metrics and neurophysiological alignment via linear regression.\n-  Apart from the stimulus synthesis method, which I believe to be well-motivated and a sensible approach to modeling the inputs, I believe the theoretical advance of this study is marginal, and am unconvinced that the proposed architecture that is central to the paper provides a significant advancement to our understanding of cortical motion processing.\n\nUPDATE: Sep 1, 2023. I have read the rebuttal, my issues with the paper remain, and I maintain my score.", "questions": "- The paragraph introducing the architecture (L108) remains opaque to me. It would be helpful to include a more thorough description of the dimensionality of x, z1, z2, z3. Are you concatenating z1,2,3 to form your 420D cNVAE latent, and comparing it to a single 420D z VAE latent? This is what I am inferring from Fig. 3.\n- There are many acronyms that are introduced and never explicitly written out. I am familiar with much of the VAE and disentanglement literature but the general reader shouldn't have to be. E.g. NVAE, DCI, TCVAE (since this one isn't pertinent to your study you can just remove it altogether instead of introducing more terms to the reader), .\n- Please elaborate on what end-point error is (L137).\n- Please include a sentence on what permutation importance is; I had to look up the sklearn docs.\n- VAEs are inherently stochastic neural networks with exact mean and covariance structure available to the experimenter. However, covariance is typically thrown out most studies using VAEs, including this one. This is at odds with the fact that noise exists and neural coding is impacted by noise (e.g. the classic Averbeck 2006 review on neural correlations).  You mention there are repeats in some trials of the CRCNS dataset; is it possible to quantify whether or not your model captures the stochastic structure of the neural responses?", "limitations": "There was discussion of limitations.", "rating": 4, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- I enjoyed the proposed stimulus synthesis method for training, and the last figure teasing apart contributions of the different generative factors to performance is nice.\n- Writing is mostly clear and easy to follow (clarity suggestions below).\n- The use of hierarchical VAEs to producing more disentangled/unentangled motion representations is interesting.", "weaknesses": "- The logic surrounding the physiological alignment analyses feels unsettlingly circular and/or redundant. This perhaps applies to Higgins' paper too since they do similar analyses. 1) The cNVAE is shown to produce improved DCI scores over vanilla VAE. 2) Your assessment of alignment effectively penalizes models with worse disentanglement  (i.e. a more distributed code), and therefore adjudicates that the cNVAE has better alignment with physiological data. \n- In fact, a central motivation of this study relies on a conjecture that disentanglement and/or unentanglement are normative goals of the dorsal stream. While this objective would be appealingly analogous to those proposed in the object recognition and ventral visual stream literature, I am unconvinced that this is the case given the evidence presented here. \n- The above points are compounded by the fact that, as you mentioned in Fig. 7, the non-hierarchical VAE architecture does comparably on predicting MT responses to the cNVAE.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "First, we want to contextualize all “weaknesses” about the brain alignment score (Fig. 8) in the greater scope of the manuscript. Fig. 8 is one of the 5 metrics we use to evaluate the learned representations (the others being untangling, disentanglement, completeness, and neural prediction). All 5 metrics favor the cNVAE. Therefore, this is a small fraction of the results.\n\n> The logic surrounding the physiological alignment...\n\nWe interpret your comment to mean our alignment metric favors disentangled codes, therefore any disentangled code will better align with MT (or any other representation, for that matter).  Although we agree that our alignment score is far from perfect, there are several reasons why we believe this is not a weakness:\n\n1. Maybe “redundant” but not “circular”: If the alignment metric was biased to favor disentangled codes, this would only make it redundant because we do not start from a motivation that codes should be disentangled (more below).\n\n2. There is no such thing as a *universally* disentangled code, and therefore this alignment score would have to depend on what is being disentangled. In its most common definition, disentanglement depends explicitly on how we define ground truth variables. We explain this point in the appendix (section: “disentanglement is in the eyes of the beholder”; Fig. S8) and show that disentanglement scores change with different ground truth definitions. \n3. The disentanglement score for the cNVAE is not that different from vanilla VAE (Fig. 5). Where the models differ most is in untangling and completeness. Given this, we think it is unlikely that disentanglement is driving the high MT alignment.\n4. Our alignment score (Fig. 8) is more dependent on the hierarchy than DCI metrics, as indicated by the good performance of the cNAE.\n5. We believe *untangling* is the most important feature of cNVAE, and where it shines (Fig. 4b). Despite their similar names, “untangling” and “disentanglement” are independent concepts, where “untangling” simply means the ground truth factors are linearly decodable ([DiCarlo & Cox, 2007](https://tinyurl.com/dicarlocox)). Although several recent lines of work have argued that disentangled codes are desirable, others have argued that distributed codes are good as long as they support untangling ([Rigotti et al., 2013](https://tinyurl.com/rigotti13)). \n6. The reviewer is holding us to a standard that no prior work is held to. When it comes to prediction performance, the best model is cNVAE (at $\\beta=0.8$) with an improvement of 0.008 over the best VAE (see Table 3). In comparison, Mineault et al. (ref 29, which was a spotlight NeurIPS paper) report a minuscule improvement of 0.001 to select their best model. \n\nWe completely agree with the reviewer that our measure of alignment is not perfect and we highlighted limitations of prediction and alignment in our main text and supplemental material. We think the reviewer is correct to be dissatisfied with disentanglement and linear regression as a measure of brain alignment. There is a lot of future work to do before the field converges on metrics for evaluating the similarity of pairs of representations, as highlighted in recent work (e.g., [Han et al. 2023](https://tinyurl.com/hanicml23)). A potential route forward is geometric analyses ([Williams et al., 2021](https://tinyurl.com/awshape21)). Ultimately, we don’t think disentanglement is central to our results, but is one of several metrics that can be used to evaluate a code. Because Fig. 8 is a small part of our results, we are happy to move it to the supplemental material and modify the text accordingly if the reviewer feels we overstated its importance.\n\n> In fact, a central motivation of this study...\n\nThat was not a motivation for this study. Our study was motivated by conjectures from Helmholtz and Mumford (see our global rebuttal). In our efforts to compress our introduction to fit the page requirements, we collapsed this idea too much and this probably misled the reviewer. We will fix this in our revisions and address this point directly here.\n\nWe were primarily motivated by the idea that hierarchical inference is important for representation learning, which we evaluated using a number of metrics that have been proposed by neuroscientists and ML researchers and concluded that hierarchical inference leads to several improvements. Our approach contrasts with previous work on dorsal stream (e.g., Mineault et al. 2021) where linearly decoding ground truth was literally the objective (they used supervised ResNets). Our models were trained solely using the standard ELBO or beta-VAE loss. Our results demonstrate that unsupervised models w/ hierarchical inference untangle more than models w/o hierarchy. Importantly, untangling and disentanglement are metrics we use to evaluate the role of hierarchy in learned representations, and they were neither a motivation nor objective.\n\n> The above points are compounded…\n\nOur best model was a cNVAE by 0.008 over the best VAE, which is 8 times larger than the selection criterion used by the SOTA model of the dorsal stream (Mineault et al., 2021), which we outperform by over a factor of 2.\n\n> VAEs are inherently stochastic…\n\nWe agree there is an interesting parallel between the covariance in VAEs and the covariance of neurons in the brain. One practical reason why we cannot examine this is our dataset consists of single-unit recordings, making it impossible to compare the covariance of these neurons to VAEs.\n\nInterestingly, we did look at the covariance of latent representations and found that cNVAE encodes its input in a much larger dimensionality than the vanilla VAE (Figure R1-f in the global rebuttal). What information is encoded in those correlations? Is it related to noise correlations? We don’t know, but we like this suggestion and plan to study this later with a population-recording dataset.\n\n> Questions\n\nWe will implement the remaining suggestions in our revisions."}]

Author rebuttal:
Here we address the most common concerns and highlight additional analyses inspired by them. We realize our work did not come across clearly to all reviewers and we offer a brief summary of our main points first:

We were motivated by the idea that representations of the natural sensory world involve the inference of the underlying causes of the senses. This is foundational in theories of perception, that posit that our brains learn hierarchical generative models of the world. We focused on the role of hierarchical inference in the learned representations of natural motion and its causes (moving objects and self-motion). Our paper’s main points are:

(i) We created a simulation framework (ROFL) for synthesizing motion stimuli. ROFL enables control over ecologically relevant factors (self-motion and objects) while avoiding confounds due to texture. Importantly, ROFL has a hierarchy of spatial scales — just like the real world — which interact in nontrivial ways (see Figs. 1c-d).

(ii) We introduced compressed NVAE (cNVAE) that greatly reduced the number of latents.

(iii) We evaluated the representations of hierarchical and non-hierarchical models using a multitude of metrics and found the cNVAE performed favorably.

(iv) We measured neural prediction and alignment of the cNVAE and comparison models using recordings from area MT, where we outperform the SOTA by a factor of 2.

In sum, we focused on understanding the relationship between the learned representations (VAE latents) and “ground truth” causes of sensor data. We showed how representations depend on architecture (hierarchical vs. non-hierarchical), loss (variational vs. maximum likelihood), and the training set.

We will add a Related Work section to the manuscript in the appendix, which should further clarify the contributions of the manuscript. As we are limited to 6K characters, we have included an abridged draft below.

# Related work (draft)
## Neuroscience and VAEs
The connection between VAEs and neuroscience is reviewed by (Marino 2022), but direct comparisons have been limited to these recent papers:

Higgins 2021: Trained beta-VAE models on face images and found that beta-VAE discovers individual latents that are aligned with IT neurons.

Storrs 2021: Trained PixelVAE on synthesized images of glossy surfaces and found that the model spontaneously learned to cluster images according to underlying physical factors and mimicked human perceptual judgments.

Csikor 2022: Investigates the properties of representations and inference in a biologically inspired hierarchical VAE called Topdown VAE that captures key properties of representations in V1 and V2 of the visual cortex.

Miliotou 2023: learned mappings from the latent space of a hierarchical VAE to fMRI voxel space supported improved reconstruction of images from brain data. Ablation experiments find hierarchy is an essential component.

## Hierarchical VAEs
Ladder VAE (LVAE) was the first to introduce hierarchy to VAEs, which improved upon standard VAEs by sharing information between the inference and generative networks, allowing LVAEs to learn deeper representations. Building on LVAE, Maaløe et al. 2019 introduced Bidirectional-Inference Variational Autoencoder (BIVA), using skip-connections to further enhance the flow of information among latent variables. Both the LVAE and BIVA enabled VAEs to effectively leverage deep stochastic hierarchies.

Recently, the hierarchical Nouveau VAE (NVAE) (Vahdat & Katz, 2020) achieved SOTA in several benchmarks and generated high-quality faces. Very Deep VAE (vdvae; Child 2021) achieved impressive performance on complex image benchmarks. Neither work evaluated how the hierarchical latent structure changed the quality of learned representations. As far as we know, ours is the first study focused on the evaluation of representations in hierarchical VAEs with applications to neuroscience data.

Additionally, NVAE and vdvae have an undesirable property: their convolutional latents result in a latent space that is several orders of magnitude larger than the input space defeating the main purpose of autoencoders. For vdvae, a tiny subset (only 3%) of its latent space is necessary [Hazami 2022](https://arxiv.org/abs/2203.13751). We demonstrate that it is possible to compress hierarchical VAEs and focus on investigating the latents.

## Evaluating DNNs on predicting biological neurons.
Many studies have evaluated DNNs for predicting brain responses, but most are on static image processing (“ventral stream”), most notably. In contrast, motion processing in the dorsal stream has only been considered thus far in Mineault (2021), who used a deep neural network to extract ground truth variables from simulated drone flight ("AirSim"). They evaluate neural prediction across many areas and achieved SOTA in prediction for the dataset we consider here, although we greatly outperform them.


# New rebuttal results, future work
## Ablation experiments
We performed ablation experiments and found that this offers insight into what is represented by the latents in cNVAE. Lesioning top latents in the encoder pathway removes the object from the reconstructed output, leaving an unperturbed background (see our rebuttal Figure R1-a). In contrast, bottom latent lesions disrupt self-motion and leave object position and velocity mostly undisturbed (Fig R1-d). We extended these ablations to our MT neuron predictions. The performance drops most dramatically when the bottom latents are disrupted. Lesioning the decoder did not lead to significant performance drops (Figure R1-e; top).

## Manifold analysis
We performed a manifold analysis using effective dimensionality (ED; see Figure R1 for a definition) and found ED is larger for cNVAE compared to VAE even though both models have the same number of latents and achieve almost identical loss (Fig. S5). This suggests a future direction to understand how cNVAE accomplishes these properties.

Meta-review: This paper examines whether hierarchical Bayesian inference could provide an account of motion processing in the primate brain. The authors develop a synthetic dataset for training networks on optical flow data, and examine the representations formed in various types of variational autoencoders. The authors show that hierarchical VAEs trained on this dataset develop representations that untangle factors of variation, are disentangled, and align well to real neuron responses in monkey MT in response to moving dot stimuli.

This was a challenging decision, as the scores were borderline. The critical reviewers were concerned that the paper failed to really provide insight into what the key factors were for model alignment, that the hierarchy didn't actually seem to be all that important for alignment to the neural data, and that the methods for assessing the model were circular. The authors provided extensive rebuttal with new data and many changes proposed. In discussion, the reviewers were then also concerned that there were too many changes promised, and the paper might change too much in the interim.

As an AC, after reading through the paper, further considering the reviews, and looking at the authors' rebuttals, I ended up coming down in favour of acceptance. The central reason I ended up deciding on accept was that the paper does present an interesting new synthetic dataset, and the fact that a hierarchical VAE trained on this dataset matches responses in primate MT to moving dot stimuli is interesting and would seem to support a theory of Bayesian inference accounting for the responses in the dorsal stream. As well, the authors did work hard to address the reviewer concerns.

However, I also want to caution the authors on a couple of points and encourage them to highlight these limitations (as well as making all the changes they promised):

1. The reviewer's concern that the non-hierarchical model also accounts for MT responses well is quite valid. That would seem to suggest that *hierarchical* inference is not actually all that important for explaining the responses, just inference. In rebuttal, the authors suggest that predictive power alone shouldn't be weighed too heavily. Which may be true, but the fact remains that predictive power isn't *meaningless*, so the authors should recognize this explicitly as a complication for their conclusions.

2. When comparing to the Mineault et al. (2021) paper, the authors recognize that one big limitation of their study relative to this previous one is the fact that they used a model that can only receive moving dot stimuli. Indeed, the R numbers they report for MT for the model in Mineault et al. (2021) are therefore only the numbers for one dataset explored in that work, but their model actually achieved a much higher R value (.381) for a video dataset. This should be explicitly mentioned in the discussion section where these issues are touched on, i.e., it's only a 2x improvement for the dots stimuli, and we can't assess the performance on the video stimuli due to the modelling decisions here. It is possible that a hierarchical VAE trained and tested on video stimuli would align more poorly than the model from that previous work.
Decision: Accept (poster)