Toward a mechanistic understanding of inference in visual cortex and diffusion models

TL;DR: Take the classic sparse coding model of V1, replace its independence assumption with a learnable pairwise interaction matrix between neurons, and train the whole thing as a denoiser with the diffusion-model objective. What you get is a minimal diffusion model that is also a recurrent neural circuit whose every parameter is inspectable. The learned lateral connections reproduce the collinear "association field" structure of V1 horizontal connections, the model denoises nearly as well as a parameter-matched U-Net, and — because the denoising Jacobian factors into an explicit gated-spreading process — you can read off, term by term, how a diffusion model assigns high probability to a continuous family of image deformations it never saw during training. Bonus surprise: about 30% of the neurons learn to disconnect from the image entirely and act as an emergent hidden layer enforcing global consistency.

Two black boxes, one problem

There is an old observation from psychophysics: scatter oriented Gabor patches randomly across an image and it looks like noise; align a subset of them along a smooth path and observers instantly perceive a contour, even though no pixels connect the elements. Neuroanatomy offers a suspect: the long-range horizontal connections in superficial V1, which preferentially link neurons with similar orientation tuning at collinearly aligned positions. The theory — Field, Hayes & Hess's "association field" — is that this wiring implements contour integration. But nobody had shown this connectivity emerging from a probabilistic generative model trained on natural images.

Meanwhile, diffusion models trained on natural images do something eerily similar. Feed the fragmented-contour image into a trained denoiser and it completes the contour. Probe the denoiser's Jacobian — which, by Tweedie's formula, is J(x) = I + \sigma^2 \nabla_x^2 \log p_\sigma(x), so its top eigenvectors point toward the model's high-probability directions — and you find oscillatory patterns hugging the image contours. Prior work (Kadkhodaie et al.) showed these "geometry-adaptive harmonic bases" are exactly what lets diffusion models generalize instead of memorizing. But the U-Net that produces them is a black box: we can see that the model learned these features, not how any circuit computes them.

This paper builds a model simple enough to answer both questions at once, because it is simultaneously a V1 circuit model and a diffusion model, and its Jacobian can be decomposed by hand.

The model: sparse coding with a learnable social network

Standard sparse coding explains an image x as x \approx \Phi z with a sparse latent code z, minimizing an energy \frac{1}{2\sigma^2}\|x - \Phi z\|_2^2 + \lambda \|z\|_1. Trained on natural images, the dictionary \Phi converges to Gabor-like receptive fields — the classic Olshausen & Field result. But the L_1 prior treats every latent as statistically independent, which is badly wrong for natural images: edges co-occur along contours, across scales, across orientations. The independence assumption is precisely what prevents standard sparse coding from doing contour completion.

The fix here is minimal and unconstrained: add a pairwise interaction term to the prior,

E_{\theta,\sigma}(x,z) = \tfrac{1}{2\sigma^2}\|x - \Phi z\|_2^2 + \gamma_\psi(\sigma)\big(\|\lambda \circ z\|_1 + \tfrac{1}{2} z^\top M z\big),

where M is a fully learnable matrix of pairwise couplings between latents (negative entries = mutual facilitation, positive = suppression, under energy minimization), \lambda is a per-neuron sparsity level, and \gamma_\psi(\sigma) is a small learned MLP that rebalances prior vs. likelihood as a function of the noise level \sigma.

Non-factorial priors like this have historically been a training nightmare under maximum likelihood — you need to marginalize a high-dimensional latent. The paper's move is to sidestep this with denoising score matching. Under a MAP approximation (assume the posterior over z is peaked at its mode z^*), the score of the marginal takes a closed form:

\nabla_x \log p_{\theta,\sigma}(x) \approx \tfrac{1}{\sigma^2}\big(\Phi z^*(x;\sigma) - x\big),

and plugging this into the DSM loss makes the noisy image cancel algebraically, leaving a plain weighted reconstruction objective: infer z^* for the noisy input, then push \Phi z^* toward the clean image. Training a structured prior reduces to training a denoiser — which is exactly what diffusion models do. The correspondence is not an analogy; the model is a (minimal, MAP-approximated) diffusion model, sampled with a standard EDM-style probability-flow ODE.

Inference for z^* runs ISTA, and the ISTA update is the neural circuit:

z_{t+1} = \text{ReLU}\Big(z_t + \eta\big[\underbrace{\Phi^\top x}_{\text{feedforward}} - \underbrace{(\Phi^\top\Phi + \sigma^2\gamma(\sigma) M)\, z_t}_{\text{recurrent}} - \underbrace{\gamma(\sigma)\lambda}_{\text{bias}}\big]\Big).

Note the crucial \sigma^2\gamma(\sigma) factor: as the input gets noisier, the recurrent prior M is weighted up relative to the bottom-up evidence. The circuit automatically trusts context more when the data is ambiguous. Training through the recurrence uses implicit differentiation at the fixed point (phantom gradients, 1–3 unrolled steps) rather than backprop through time.

The ISTA fixed-point iteration is the circuit: bottom-up drive Φᵀx, lateral recurrence through the Gram matrix plus the learned interaction matrix M, and a sparsity threshold. The σ²γ(σ) gain means the learned prior dominates exactly when the sensory evidence is weak — trained end-to-end with the diffusion (denoising score matching) objective.

What it learns: an association field, from natural images

Trained on grayscale natural scenes (multi-scale convolutional implementation, Appendix A.8), \Phi converges to oriented Gabors as usual. The interesting object is M. Its excitatory entries connect spatially offset neurons whose preferred orientations co-align along a shared axis — collinear facilitation — alongside strong self-inhibition enforcing sparsity. Compared against the Gram matrix \Phi^\top\Phi (which only mediates local, overlapping-receptive-field interactions), the learned excitatory connections in M span systematically larger spatial displacements and link units of similar preferred orientation (Fig. S2). This is qualitatively the horizontal-connection structure documented in superficial V1, emerging from a generic denoising objective on natural images — something previous models largely hand-built into the prior.

The functional payoff shows up against an ablated control with M = 0 (i.e., ordinary factorial sparse coding trained the same way): the factorial model hallucinates dense spurious edges in noisy regions because it evaluates each feature in isolation, while the non-factorial model completes true contours and suppresses uncoordinated clutter (Fig. 2A). Against a parameter-matched U-Net (attention disabled, ~4 resolution levels), the model reaches comparable denoising PSNR across noise levels and produces visually similar reverse-diffusion samples (Fig. 6) — the "generative parity" claim.

The aha: reading the Jacobian like a wiring diagram

Here is the part that justifies the paper. Because the denoised output is linear in the fixed point, D_\theta(x) = \Phi z^*, the pixel-space Jacobian factors exactly:

J(x) = \Phi\, J_z\, \Phi^\top, \qquad J_z = \eta\big[I - \Sigma(I - \eta W)\big]^{-1}\Sigma,

where W = \Phi^\top\Phi + \sigma^2\gamma(\sigma)M collects all lateral interactions and \Sigma = \mathrm{diag}(\mathbf{1}_{z^* > 0}) is a binary mask of which neurons are currently active. Expand the inverse as a Neumann series and look at one column — the network's impulse response to nudging a single active neuron i:

J_{z,i} = e_i + \eta\,\Sigma W e_i + \eta^2\,\Sigma W \Sigma W e_i + \dots

Every term is interpretable. \Sigma instantly kills any perturbation to inactive neurons — only contextually relevant signals survive. W spreads the surviving signal along the learned lateral connections, with a collinear bias. Then \Sigma clamps again: activity may only leak into neurons that are already active. The k-th term is k-hop spreading, ricocheting through the network but strictly confined to the active support — in practice, the active contour.

One column of the latent Jacobian, unrolled as a Neumann series. Unconstrained spreading through W has a collinear bias but leaks off-axis (small stray needles). Alternating spread with the sparse gate Σ — spread, mask, spread, mask — confines the impulse response to the currently active contour. Summing over hops yields eigenvectors that are oscillatory modes living exactly on the image's contours.

This is the mechanistic answer to the diffusion-generalization question. The dominant Jacobian eigenvectors are not stored templates; they are emergent oscillatory modes of a gated lateral-spreading process, adapted on the fly to whatever contours are active in the input. Because they span a continuous subspace of high-probability directions, any linear combination is another valid image state: the model can shift and bend contours while preserving their integrity, assigning high probability to infinitely many deformations of finitely many training examples. Kadkhodaie et al. showed empirically that black-box diffusion models produce exactly such geometry-adaptive bases; this paper shows a concrete circuit that constructs them, and argues (plausibly but without direct evidence) that residual U-Nets are doing an amortized version of the same recurrent computation.

Detached neurons: a hierarchy nobody asked for

Trained on CelebA, both the U-Net and this model produce Jacobian eigenvectors encoding globally coordinated deformations — both eyes shifting together, face geometry deforming coherently. How does a single-layer circuit coordinate across the whole image? Inspecting \Phi, the authors find that over 30% of the latent units learn near-zero dictionary norms: they neither receive meaningful visual input nor contribute pixels to the reconstruction. These "detached neurons" interact with the rest of the network only through M.

They are not dead weight. Ablating their recurrent interactions causes a catastrophic drop in denoising PSNR at high noise — exactly the regime where the network must hallucinate structure from prior knowledge — and the Jacobian eigenvectors lose their long-range coordination, degenerating into disconnected local deformations. The interpretation: the flat architecture spontaneously carved out a hidden second layer that enforces global consistency through lateral connections alone. For anyone thinking about why depth and hierarchy emerge, this is a striking result — hierarchy as an attractor of the objective, not an architectural choice.

Predictions for real V1

The framework generates concrete physiology hypotheses (Section 4.4). Because the recurrent drive is \Phi^\top\Phi + \sigma^2\gamma(\sigma)M, the circuit should shift from likelihood-dominated to prior-dominated dynamics as stimulus ambiguity rises. At high contrast, perturbing a neuron should suppress iso-oriented neighbors — consistent with recent optogenetic findings (Chettih & Harvey-style single-cell perturbations). At low contrast, the model predicts the sign flips: excitatory, collinear-specific surround facilitation, gated by the active population \Sigma so that the spread follows the stimulus contour when it aligns with the target's preferred orientation and is extinguished when orthogonal. Contrast-dependent surround sign-flips are already documented in V1; the orientation-and-context-specific gating structure is the new, testable part.

How strong is the evidence?

The mechanism is airtight in a way most interpretability work is not: Equations 7–8 are exact consequences of the model, not post-hoc probes. The soft spots are elsewhere.

The V1 correspondence is qualitative. Figure 2B and the Fig. S2 statistics (longer-range excitation in M than in the Gram matrix, orientation similarity of top connections) match the shape of the association field, but there is no quantitative fit to actual anatomical or functional connectivity data — no comparison to Bosking-style tracer statistics or measured co-circularity kernels. "Mirrors the structure of horizontal connections" is fair; "matches V1" would not be.

"Generative parity" rests on PSNR curves and visual sample comparison against one parameter-matched, attention-free U-Net at 64×64 on CelebA and small natural-image patches. No FID, no larger baselines, no diversity metrics. At this scale the claim is credible; whether the mechanism story scales to what a modern DiT does with attention is entirely open — the paper's discussion gestures at residual networks performing implicit recurrence, citing others rather than demonstrating it. The MAP approximation of the score is also uncontrolled: the model matches an approximate score, so it is a biased diffusion model, and the paper does not quantify the gap. And the detached-neuron result, while fascinating, gets a mostly qualitative treatment — one dataset, one ablation style.

None of this undermines the core contribution. This is the cleanest existing bridge between a documented cortical circuit motif and the inference machinery of diffusion models, and it is fully white-box. If the picture holds, it reframes diffusion generalization as gated lateral spreading on a sparse support — a hypothesis you can now go test inside U-Nets — and it hands V1 physiologists a specific, noise-contingent perturbation experiment.

Where to spend your time: Section 3.4 plus Figure 3 (the Jacobian decomposition — the paper's heart), then Section 4.3 (detached neurons). Appendix A.8 is worth a skim if you want to implement it: the real system is multi-scale convolutional with block-tridiagonal cross-scale interactions and Gauss–Seidel sweeps, considerably more engineered than the clean single-scale equations suggest.