SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Pawel A. Pierzchlewicz*,1-2, Konstantin F. Willeke1-2, Arne F. Nix1-2, Pavithra Elumalai2, Kelli Restivo3-4, Tori Shinn3-4, Cate Nealley3-4, Gabrielle Rodriguez3-4, Saumil Patel3-4, Katrin Franke3-4, Andreas S. Tolias3-5, Fabian H. Sinz1-4
$^{1}$ Institute for Bioinformatics and Medical Informatics, Tübingen University, Tübingen, Germany  $^{2}$ Institute of Computer Science and Campus Institute Data Science, University of Göttingen, Germany
$^{3}$ Department of Neuroscience, Baylor College of Medicine, Houston, TX, USA
$^{4}$ Center for Neuroscience and Artificial Intelligence, Baylor College of Medicine, Houston, TX, USA
$^{5}$ Department of Electrical and Computer Engineering, Rice University, Houston, TX, USA
*ppierzc@cs.uni-goettingen.de
In recent years, most exciting inputs (MEIs) synthesized from encoding models of neuronal activity have become an established method for studying tuning properties of biological and artificial visual systems. However, as we move up the visual hierarchy, the complexity of neuronal computations increases. Consequently, it becomes more challenging to model neuronal activity, requiring more complex models. In this study, we introduce a novel readout architecture inspired by the mechanism of visual attention. This new architecture, which we call attention readout, together with a data-driven convolutional core outperforms previous taskdriven models in predicting the activity of neurons in macaque area V4. However, as our predictive network becomes deeper and more complex, synthesizing MEIs via straightforward gradient ascent (GA) can struggle to produce qualitatively good results and overfit to idiosyncrasies of a more complex model, potentially decreasing the MEI's model-to-brain transferability. To solve this problem, we propose a diffusion-based method for generating MEIs via Energy Guidance (EGG). We show that for models of macaque V4, EGG generates single neuron MEIs that generalize better across varying model architectures than the state-of-the-art GA, while at the same time reducing computational costs by a factor of  $4.7\mathrm{x}$ , facilitating experimentally challenging closed-loop experiments. Furthermore, EGG diffusion can be used to generate other neurally exciting images, like most exciting naturalistic images that are on par with a selection of highly activating natural images, or image reconstructions that generalize better across architectures. Finally, EGG is simple to implement, requires no retraining of the diffusion model, and can easily be generalized to provide other characterizations of the visual system, such as invariances. Thus, EGG provides a general and flexible framework to study the coding properties of the visual system in the context of natural images. $^{1}$
From the early works of Hubel and Wiesel [1], visual neuroscience has used the preferred stimuli of visual neurons to gain insight into the information processing in the brain. In recent years, deep learning has made big strides in predicting neuronal responses [2-16] enabling in silico stimulus
<sup>1</sup>The code is available at https://github.com/sinzlab/energy-guided-diffusion
Figure 1: Schematic of the EGG diffusion method with a pre-trained diffusion model. Examples of applications: Left: Most Exciting Inputs for different neurons, Middle: Most Exciting Naturalistic Inputs matched unit-wise to the MEIs. Right: Reconstructions in comparison to the ground truth (top) and gradient descent optimized (bottom).
synthesis of non-parametric most exciting inputs (MEIs) [17-19]. MEIs are images that strongly drive a selected neuron and can thus provide insights into its tuning properties. Up until now, they have been successfully used to find novel properties of neurons in various brain areas in mice and macaques [17-24].
However, as we move up the visual hierarchy, such as monkey visual area V4 and IT, the increasing non-linearity of neuronal responses with respect to the visual stimulus makes it more challenging to obtain models with high predictive performance for single neurons, and optimize perceptually plausible MEIs, that is, those not corrupted by adversarial high-frequency noise for example. Particularly, area V4 is known to be influenced by attention effects [25], and shifts in attention before the onset of saccades can change the location of its neurons' receptive fields [26, 27]. When models become more complex or units are taken from deeper layers of a network, existing MEI optimization methods based on gradient ascent (GA) can sometimes have difficulties producing qualitatively good results [28] and can overfit to the idiosyncrasies of more complex models, potentially decreasing the MEI's model-to-brain transferability. Typically, these challenges are addressed by biasing MEIs towards the statistic of natural images, for instance by gradient pre-conditioning [28], by including a total variation loss to reduce high-frequency noise [29] or by image synthesis via GANs [19]. However, as discussed by Engstrom et al. [30] and Feather et al. [31] including additional priors into the generation process can result in obfuscated model biases.
Here, we make two contributions towards the above points: 1 We introduce a new model architecture, called the attention readout, for predicting the activity of neurons in macaque area V4 , which together with a data-driven convolutional core outperforms previous task-driven models [24, 32]. 2 To improve the quality of MEI synthesis we introduce a novel method for optimizing MEIs via Energy Guided Diffusion (EGG). EGG diffusion guides a pre-trained diffusion model with a learned neuronal encoding model to generate MEIs with a bias towards natural image statistics. Our proposed EGG method is simple to implement and, in contrast to similar approaches [33-35], requires no retraining of the diffusion model (Fig. 1). We show that EGG diffusion not only yields MEIs that generalize better across architectures and are thus expected to drive real neurons equally well or better than GA-based MEIs but also provides a significant  $(4.7\mathrm{x})$  speed up over the standard GA method enhancing its utility for close-loop experiments such as inception loops [17, 18, 20, 24]. Since optimizing MEIs for thousands of neurons can take weeks [24], such a speed-up directly decreases the energy footprint of this technique. Moreover, the rapid verification of synthesized images in vivo is particularly important for close-loop experiments given that maintaining the stability of single unit recordings is challenging, and there's also the issue of representational drift [36], where tuning functions can change over time. We also demonstrate that EGG diffusion straightforwardly generalizes to provide other characterizations of the visual system that can be phrased as an inverse problem, such as image reconstructions based on neuronal responses. The flexibility and generality of EGG thus make it a powerful tool for investigating the neural mechanisms underlying visual processing.
2 Attention readout for macaque area V4
Background Deep network-based encoding models have set new standards in predicting neuronal responses to natural images [2-15]. Virtually all architectures of these encoding models consist of at least two parts: a core and a readout. The core is usually implemented via a convolutional
network that extracts non-linear features  $\Phi(x)$  from the visual input and is shared across all neurons to be predicted. It is usually trained through one of two paradigms: i) task-driven, where the core is pre-trained on a different task like object recognition [3, 4, 37-39] and then only the readout is trained to predict the neurons' responses or ii) data-driven where the model is trained end-to-end to predict the neurons' responses. The readout is a collection of predictors that map the core's features to responses of individual neurons. With a few exceptions [40], the readout components and its parameters are neuron-specific and are therefore kept simple. Typically, the readout is implemented by a linear layer with a rectifying non-linearity. Different readouts differ by the constraints they put on the linear layer to reduce the number of parameters [3, 4, 37, 40-42]. One key assumption all current readout designs make is that the readout mechanism does not change with the stimulus. In particular, this means that the location of the receptive field is fixed. While this assumption is reasonable for early visual areas like V1, it is not necessarily true for higher or mid-level areas such as macaque V4, which are known to be affected by attention effects and can even shift the location of the receptive fields [26]. This motivated us to create a more flexible readout mechanism for V4.
State-of-the-art model: Robust ResNet core with Gaussian readout In this study, we compare our data-driven model to a task-driven model [24], which is also composed of a core and readout. The core is a pre-trained robust ResNet50 ( $L_{2}, \varepsilon = 0.1$ ) [43, 44]. We use the layers up to layer 3 in the ResNet, which has 1,024 channels, thus providing a 1,024 dimension feature space. Then batch normalization is applied [45], followed by a ReLU non-linearity. The Gaussian readout [40] learns the position of each neuron and extracts a feature vector at this position. During training, the positions are sampled from a 2D Gaussian distribution with means  $\mu_{n}$  and  $\Sigma_{n}$ , during inference the  $\mu_{n}$  positions are used. Then the extracted features are used in a linear non-linear model to predict neuronal responses. We will refer to this model as the Gaussian model.
Proposed model: Data-driven core with attention readout The predictive model is trained from scratch to predict the neuronal responses in an end-to-end fashion. Following Lurz et al. [40], the architecture is comprised of two main components. First, the core, a four-layer CNN with 64 channels per layer with an architecture identical to Lurz et al. [40]. Secondly, the attention readout, which builds upon the attention mechanism [46, 47] as it is used in the popular transformer architecture [48]. After adding a fixed positional embedding to  $\Phi (\pmb {x})$  and normalization through LayerNorm [49] to get  $\tilde{\Phi} (\pmb {x})$ , key and value embeddings are extracted from the core representation. This is done by position-wise linear projections  $V\in \mathbb{R}^{c\times d_k}$  and  $U\in \mathbb{R}^{c\times d_v}$  both of which have parameters shared across all neurons. Then, for each neuron a learned query vector  $\mathbf{q}_n\in \mathbb{R}^{d_k}$  is compared with each position's key embedding using scaled dot-product attention [48].
$$
\alpha_ {n} = \operatorname {s o f t m a x} \left(\sum_ {c, d _ {k}} \frac {\tilde {\Phi} (\boldsymbol {x}) _ {c} W _ {c , d _ {k}} q _ {n , d _ {k}}}{\sqrt {d _ {k}}}\right) \tag {1}
$$
The result is a spatially normalized attention map  $\alpha_{n}\in \mathbb{R}^{h\times w\times 1}$  that indicates the most important feature locations for a neuron  $n$  given an input image. Using this attention map to compute a weighted sum of the value embeddings gives us a single feature vector for each neuron. Finally, a neuron-specific affine projection with ELU non-linearity [50] gives rise to the predicted spike rate  $\hat{r}_n$  (Fig. 2A). The model training is performed by minimizing the Poisson loss using the same setup as described in Willeke et al. [24]. We will refer to this model as the Attention model.
Training data We use data from 1,244 Macaque V4 neurons from Willeke et al. [24] and briefly summarize their data acquisition in the supplementary materials section A.1.
Results Our Attention model significantly outperforms the Gaussian model in predicting neuronal responses of macaque V4 cells on unseen natural and model-derived images. We evaluate the model performance by the correlation between the model's prediction and the averages of actual neuron responses across multiple presentations of a set of test images, as described by Willeke et al. [24]. We compared this predictive performance to the Gaussian model [44] on 1,244 individual neurons (Fig. 2B). The Attention model significantly outperforms the Gaussian model by  $12\%$  (Wilcoxon signed-rank test,  $\mathrm{p - value} = 6.79\cdot 10^{-82}$ ). In addition, we evaluated the new readout on how well it predicts the real neuronal responses to 48 MEIs generated from the Gaussian model [see 24] and 7 control natural images. Our Attention model is better at predicting real neuronal responses,
Figure 2: a) Schematic of the Attention Readout. b) Correlation to average scores for 1,244 neurons. The Attention model (pink) shows a significant (as per the Wilcoxon signed rank test, p-value  $= 6.79 \cdot 10^{-82}$ ) increase in the mean correlation to average in comparison to the Gaussian model (blue). c) Predictive performance comparison of the two models in a closed-loop MEI evaluation setting. Showing that the data-driven with attention readout model better predicts the in-vivo responses of the MEIs.
even for MEIs of another architecture (Fig. 2C). Please note that Willeke et al. [24] experimentally verified MEIs in only a subset of neurons and only used the neurons with high functional consistency across different experimental sessions. For that reason, we too can only compare the performance of model-derived MEIs on this subset of neurons. We additionally show that the Attention model and Gaussian model show representational similarity (see Table S1) and that the Attention model uses its ability to shift its receptive field (Fig. S1).
Readout \ Core	Task-Driven	Data-Driven
Factorized	-	0.153
Gaussian	0.262	0.229
Attention	0.276	0.294
Table 1: Ablation study test correlation comparison for combinations of different cores and readouts. Bold indicates the best-performing model.
Ablation Study We perform an ablation study comparing the effects of the choice of core and readout on the performance in terms of test correlation (Table 1). We identify that the data-driven core + Attention readout model outperforms all previous setups. Furthermore, the ablation study shows that the Attention readout generally improves performance across cores.
3 Energy guided diffusion (EGG)
3.1 Algorithm and methods
In this section, we describe our approach to extract tuning properties of neuronal encoding models using a natural image prior as described by a diffusion model. In brief, we use previously established links between diffusion and score-based models and the fact that many tuning properties can be described as inverse problems (most exciting image, image reconstruction from neuronal activity, etc.) to combine an energy landscape defined by the neuronal encoding model with the energy landscape defined by the diffusion model and synthesize images via energy minimization. We show that this method leads to better generalization of MEIs and image reconstructions across architectures, faster generation, and allows for generating natural-looking stimuli.
Background: diffusion models Recently, Denoising Diffusion Probabilistic Models (DDPMs) have proved to be successful at generating high-quality images [33, 51-56]. These models can be formalized as a variational autoencoder with a fixed encoder  $\pmb{x}_0 \mapsto \pmb{x}_T$  that turns a clean sample  $\pmb{x}_0$  into a noisy one  $\pmb{x}_T$  by repeated addition of Gaussian noise, and a learned decoder  $\pmb{x}_T \mapsto \pmb{x}_0$  [33],
which is often described as inverting a diffusion process [51]. After training, the sampling process is initialized with a standard Normal sample  $\pmb{x}_T \sim \mathcal{N}(\mathbf{0},\mathbf{I})$  which is iteratively "denoised" for  $T$  steps until  $\pmb{x}_0$  is reached. In the encoding, each step  $t$  corresponds to a particular noise level such that
$$
\boldsymbol {x} _ {t} = \sqrt {\bar {\alpha} _ {t}} \boldsymbol {x} _ {0} + \sqrt {1 - \bar {\alpha} _ {t}} \boldsymbol {\varepsilon} _ {0} \tag {2}
$$
where  $\bar{\alpha}_t$  controls the signal strength at time  $t$  and  $\varepsilon_0\sim \mathcal{N}(\mathbf{0},\mathbf{I})$  is independent Gaussian noise. In the decoding step, the diffusion model predicts the noise component  $\varepsilon_{\theta}(\pmb{x}_{t},t)$  at each step  $t$  of the diffusion process [33]. Then the sampling is performed according to
$$
\boldsymbol {x} _ {t - 1} = \frac {1}{\sqrt {\alpha_ {t}}} \left(\boldsymbol {x} _ {t} - \frac {1 - \alpha_ {t}}{\sqrt {1 - \bar {\alpha} _ {t}}} \varepsilon_ {\theta} \left(\boldsymbol {x} _ {t}, t\right)\right) + \sigma_ {t} \boldsymbol {z} \tag {3}
$$
where  $z\sim \mathcal{N}(\mathbf{0},\mathbf{I})$
Several previous works have established a link between diffusion models and energy-based models [57-59]. In particular, the diffusion model  $\varepsilon_{\theta}(\boldsymbol{x}_t,t)$  can be interpreted as a score function, i.e. the gradient of a log-density or energy w.r.t. the data  $\nabla_{\boldsymbol{x}}\log p(\boldsymbol{x})$  [60]. This link is particularly useful since combining two density models via a product is equivalent to adding their score functions.
EnerGy Guided Diffusion (EGG) To optimize neurally exciting images, we require a method that can guide diffusion models via neural encoding models. The parameterization of diffusion models introduced by Ho et al. [33] only allows for the unconditioned generation of samples. Dhariwal and Nichol [53] introduced a method for sampling from a conditional distribution  $p_t(\boldsymbol{x} \mid \boldsymbol{y})$ , with diffusion models using a classifier  $p_t(\boldsymbol{y} \mid \boldsymbol{x})$  known as classifier guidance. However, this method requires i) the classifier to be trained on the noisy images, and ii) is limited to conditions for which classification makes sense. Essentially, this method relies on computing the score of the posterior distribution.
$$
\nabla_ {\boldsymbol {x} _ {t}} \log p (\boldsymbol {x} _ {t} \mid \boldsymbol {y}) = \nabla_ {\boldsymbol {x} _ {t}} \log p (\boldsymbol {x} _ {t}) + \nabla_ {\boldsymbol {x} _ {t}} \log p (\boldsymbol {y} \mid \boldsymbol {x} _ {t}) \tag {4}
$$
For classifier-guidance, the gradient of a model  $\nabla_{\boldsymbol{x}_t}\log p(\boldsymbol{y}|\boldsymbol{x}_t)$  with respect to the noisy input  $\boldsymbol{x}_t$  is combined with the diffusion model  $\nabla_{\boldsymbol{x}_t}\log p(\boldsymbol{x}_t)$ , resulting in samples  $\boldsymbol{x}_0$  conditioned on the class  $\boldsymbol{y}$ . Note that this requires a model  $\nabla_{\boldsymbol{x}_t}\log p(\boldsymbol{y}|\boldsymbol{x}_t)$  that has been trained on noisy samples of the diffusion before. Here we extend this approach to i) use neuronal encoding models, such as the ones described above, to guide the diffusion process and ii) use a model trained on clean samples only. We achieve i) by defining conditioning as a sum of energies. Specifically, we redefine equation (4) in terms of the output of the diffusion model  $\varepsilon_{\theta}(\boldsymbol{x}_t,t)$  and an arbitrary energy function  $E(\boldsymbol{x}_t,t)$ :
$$
\bar {\varepsilon} (\boldsymbol {x} _ {t}, t) = \varepsilon_ {\theta} (\boldsymbol {x} _ {t}, t) + \lambda_ {t} \nabla_ {\boldsymbol {x} _ {t}} E (\boldsymbol {x} _ {t}, t) \tag {5}
$$
where  $\lambda_{t}$  is the energy scale. This takes advantage of the fact that sampling in DDPMs is functionally equivalent to Langevin dynamics [51]. Langevin dynamics generally define the movement of particles in an energy field and in the special case when  $E(x) = -\log p(x)$ , Langevin dynamics generates samples from  $p(x)$ . For this study, we use a constant value of  $\lambda$  and normalize the gradient of the energy function to a magnitude of 1.
To achieve ii) we use an approximate clean sample  $\bar{x}_0$ , i.e. the original image, that can be estimated at each time step  $t$ . This is achieved by a simple trick introduced in Li et al. [61]. By inverting the forward diffusion process, with the assumption that the predicted  $\varepsilon_{\theta}(\boldsymbol{x}_t,t)$  is the true noise:
$$
\bar {\boldsymbol {x}} _ {0} \left(\boldsymbol {x} _ {t}, t\right) = \frac {1}{\sqrt {\bar {\alpha} _ {t}}} \left(\boldsymbol {x} _ {t} - \sqrt {1 - \bar {\alpha} _ {t}} \varepsilon_ {\theta} \left(\boldsymbol {x} _ {t}, t\right)\right). \tag {6}
$$
As a result, the energy function receives inputs that are in the domain of  $\pmb{x}_0$  at much earlier time steps  $t$ , and hence makes it feasible to use energy functions only defined on  $\pmb{x}_0$  and not  $\pmb{x}_t$ , dropping the requirement to provide an energy  $E(\pmb{x}_t, t)$  that can take noisy images. Thus, the new score can be defined as
$$
\bar {\varepsilon} (\boldsymbol {x} _ {t}, t) = \varepsilon_ {\theta} (\boldsymbol {x} _ {t}, t) + \lambda_ {t} \nabla_ {\boldsymbol {x} _ {t}} E (\bar {\boldsymbol {x}} _ {0} (\boldsymbol {x} _ {t}, t)) \tag {7}
$$
This is particularly relevant in the domain of neural system identification, as encoding models are trained on neuronal responses to natural "clean" images [2-15, 17, 21, 24, 40]. To get an energy that can understand noisy images would require showing the noisy images to the animals in experiments, which would make the use of this method prohibitively more difficult. Therefore, a guidance method that does not require training an additional model on noisy images allows researchers to apply EGG diffusion directly to existing models trained on neuronal responses and extract tuning properties from them.
Related work Many other methods have been proposed to condition the samples of diffusion processes on additional information. Ho and Salimans [55] provided a method that addressed the second requirement of classifier-guidance by incorporating the condition  $\pmb{y}$  into the denoiser  $\varepsilon_{\theta}(\pmb{x}_t, t, \pmb{y})$ . However, to introduce a conditioning domain  $\pmb{y}$  in this classifier-free guidance, the whole diffusion model needs to be retrained. Furthermore, this link between diffusion models and energy-based models allowed several previous works to compose diffusion models to generate outputs that contain multiple desired aspects of a generated image [57-59]. However, these studies focus solely on generalizing the classifier-free guidance to allow guiding diffusion models with other diffusion models. Nichol and Dhariwal [52] have used a similar gradient conditioning to guide the diffusion process using the gradient of the dot product of the CLIP image and text vectors. It has been shown that CLIP models that have not been trained on noisy images can be used for guiding diffusion models [62, 63]. Kadkhodaie and Simoncelli [64] introduced a stochastic coarse-to-fine gradient ascent procedure for generating samples from the implicit prior embedded within a CNN. While we were working on this project, Feng et al. [65] published a preprint where they used the score-based definition of diffusion models to introduce an image-based prior for inverse problems where the posterior score function is available. This work is most closely related to our approach. However, they focus on how to obtain samples and likelihoods from the true posterior. For that reason, they need guiding models to be proper score functions. We do not need that constraint and focus on guiding inverse problems defined by a more general energy function and focus particularly on the application to neuronal encoding models.
Image preprocessing for neural models The neural models used in this study expect  $100 \times 100$  images in grayscale. However, the output of the ImageNet pre-trained Ablated Diffusion Model (ADM) [53] is a  $256 \times 256$  RGB image. We, therefore, use an additional compatibility step that performs i) downsampling from  $256 \times 256 \rightarrow 100 \times 100$  with bilinear interpolation and ii) takes the mean across color channels providing the grayscale image. Each of these preprocessing steps is differentiable and is thus used end-to-end when generating the image.
3.2 Experiments