TAVAE: A VAE with Adaptable Priors Explains Contextual Modulation in the Visual Cortex

TL;DR

The paper addresses how task-specific contextual priors can be learned and deployed in the primary visual cortex. It introduces the Task-Amortized VAE (TAVAE), which preserves the VAE likelihood $p({\bm{x}}|{\bm{z}})$ while adapting the latent prior to a task-specific $p_T({\bm{z}})$, implemented as a zero-mean Laplace prior with self-consistent scale updates. Empirically, TAVAE accounts for sharpening of population responses, baseline suppression, and bimodal posterior-like patterns when task statistics are violated, and it captures how priors are updated across days in mice performing orientation discrimination. The work demonstrates that flexible, on-demand contextual priors can be learned without retraining the entire model, offering a normative framework for contextual modulation in early visual processing and a general tool for studying cortical inference under realistic task conditions.

Abstract

The brain interprets visual information through learned regularities, a computation formalized as probabilistic inference under a prior. The visual cortex establishes priors for this inference, some delivered through established top-down connections that inform low-level cortices about statistics represented at higher levels in the cortical hierarchy. While evidence shows that adaptation leads to priors reflecting the structure of natural images, it remains unclear whether similar priors can be flexibly acquired when learning a specific task. To investigate this, we built a generative model of V1 optimized for a simple discrimination task and analyzed it together with large-scale recordings from mice performing an analogous task. In line with recent approaches, we assumed that neuronal activity in V1 corresponds to latent posteriors in the generative model, enabling investigation of task-related priors in neuronal responses. To obtain a flexible test bed, we extended the VAE formalism so that a task can be acquired efficiently by reusing previously learned representations. Task-specific priors learned by this Task-Amortized VAE were used to investigate biases in mice and model when presenting stimuli that violated trained task statistics. Mismatch between learned task statistics and incoming sensory evidence produced signatures of uncertainty in stimulus category in the TAVAE posterior, reflecting properties of bimodal response profiles in V1 recordings. The task-optimized generative model accounted for key characteristics of V1 population activity, including within-day updates to population responses. Our results confirm that flexible task-specific contextual priors can be learned on demand by the visual system and deployed as early as the entry level of visual cortex.

Figures (12)

• Figure 1: a, Experimental setup. b, Go-NoGo visual discrimination task stimuli with training and testing schedule. c, Population response map of the recorded V1 population on D1 for the Go stimulus. d, Population response profile averaged over time for trained and untrained (naive) mice. e, Cartoon of VAE (top) and TAVAE (bottom) models. The TAVAE generative model component of VAE (pale orange) is reused in TAVAE whereas the recognition model of VAE (dark orange) is transformed through reweighing the variational posterior with the task-prior (Eq. \ref{['eq:p_T_conditional']}). f, Illustration of the posterior of two latent dimensions; inset: receptive fields of latents. g, Task prior of TAVAE for the discrimination task. Scale of the Laplace prior is shown. h, TAVAE responses to the Go stimulus with task (green) and natural priors (blue).
• Figure 2: Effect of learning the discrimination task in the model and in the experiment.a, b, Response profiles to task stimuli (45° and 135° gratings, red and green dashed lines, respectively) using natural prior (blue) and the task prior (solid green and red for Go and NoGo respectively) in the model, smoothing 3°. Dashed grey lines correspond to the learned contextual prior. Shading: bootstrap estimate of $95\%$ confidence interval of the mean across model neurons. c, Width of the 45° peak using the natural prior (blue) and task prior (orange; statistics are computed from different grating phases). d, Baseline reduction in TAVAE. Baseline activity was computed for each stimulus configuration from Fig. \ref{['fig:D1']}a,b and Fig. \ref{['fig:bimodal']}b. Box plot and significance markers are based on statistics computed across these stimuli. e, f, D1 Go and NoGo responses in naive (blue) and trained (green, red) mice, smoothing 6°, shades 2 sem. g, Width of the Go response over days and timeframes for naive (blue) vs. trained animals (orange, mean over neurons, trials). h, Baseline reduction following training in the experiment. Baseline activity was computed for each stimulus configuration from Fig. \ref{['fig:D1']}e,f and Fig. \ref{['fig:bimodal']}a. Boxplots: 25-75%, whiskers 1.5$\cdot$IQR. Asterisks indicate statistical significance.
• Figure 3: Systematic biases for mismatched training and test in experiment and model. a, V1 responses over five experimental sessions during which the NoGo stimulus (red) is deviating from the training NoGo stimulus at 135°, confidence shades 2 sem. Responses from naive animals are shown for reference (blue). Bimodal peaks-to-middle-trough activity differences (vertical black arrows). Mean absolute difference between stimulus and peak orientation (horizontal red and blue mirrored half arrows). b, Model responses on the same stimulus orientations as in the experiment across the five sessions. A prior learned for 45° and 90° gratings is assumed (grey dashed lines), bimodal peaks-to-middle-trough without smoothing (vertical black arrows). All panels: data points, smoothing and confidence intervals as in Fig. \ref{['fig:D1']}, asterisks indicate statistical significance (see text).
• Figure 4:
• Figure 5: a, Receptive fields for example model neurons in EAVAE base model. b Tuning curves and von Mises fits for these neurons. Note that the last neuron is classified as not orientation selective.