A Flow Matching Framework for Neural Representational Dissimilarity

Zeyuan Ye; Xue-Xin Wei

cs.AISep 25, 2026

A Flow Matching Framework for Neural Representational Dissimilarity

Authors: Zeyuan Ye, Xue-Xin Wei

Organizations: Department of Neuroscience, The University of Texas at Austin

Abstract

Neural representational dissimilarity quantifies differences between neural response distributions, and is essential for comparing neural codes across stimuli, brain areas, tasks, and models. Commonly used distance metrics involve different assumptions and are estimated with separate methods. Here, we show that a variety of distance metrics can be unified under a flow matching framework developed in deep generative models. That is, these distances arise as Jeffreys divergences under different velocity constraints. We find that flow matching has advantages for estimating distances involving complicated distributions and continuous variables. Furthermore, this framework enables the design of new distance metrics in a principled way. Together, flow matching provides a unified approach for understanding, estimating, and designing neural representational dissimilarity metrics.

Figures & tables

Figure 1: Overview of the proposed flow matching framework for neural representational dissimilarity. ( A ) Different stimuli evoke condition-specific neural responses, producing different population-response distributions whose dissimilarity is to be quantified. ( B ) Flow matching learns velocity fields

v_{1}^{\mathcal{V}}(x,t)

and

v_{2}^{\mathcal{V}}(x,t)

that transport a shared source distribution to the fitted distributions

q_{1}^{\mathcal{V}}

and

q_{2}^{\mathcal{V}}

. Here, we define representational dissimilarity as the Jeffreys divergence between

q_{1}^{\mathcal{V}}

and

q_{2}^{\mathcal{V}}

. We prove that different choices of

\mathcal{V}

correspond to different distance metrics.

Figure 2: Velocity classes $\mathcal{V}$ used in flow matching determine the fitted distributions $q_{c}^{\mathcal{V}}$ (also see Table 1 ). (A) A two-component Gaussian-mixture target, shown as samples and dotted density contours. (B) A constant velocity translates the Gaussian source (black region) without changing its shape. (C) An affine velocity allows an affine transformation of the source distribution, yielding a Gaussian approximation to the target. (D) An unconstrained velocity captures the two modes of the target. Arrows indicate velocity direction, not magnitude.

Velocity field	Induced distance
$v_{c}(x,t)=\dot{\beta}_{t}{\color[rgb]{0.5703,0.1367,0.1367}b_{c}}$	Euclidean distance
$v_{c}(x,t)=\dot{\beta}_{t}{\color[rgb]{0.5703,0.1367,0.1367}b_{c}},\ \\|{\color[rgb]{0.5703,0.1367,0.1367}b_{c}}\\|=r$	Cosine distance
$v_{c}(x,t)=\dot{\beta}_{t}{\color[rgb]{0.5703,0.1367,0.1367}b_{c}},\ \mathbf{1}^{\top}{\color[rgb]{0.5703,0.1367,0.1367}b_{c}}=0,\ \\|{\color[rgb]{0.5703,0.1367,0.1367}b_{c}}\\|=r$	Correlation distance
$v_{c}(x,t)=\dot{\beta}_{t}{\color[rgb]{0.5703,0.1367,0.1367}b_{c}}+{\color[rgb]{0.5703,0.1367,0.1367}A_{t}}(x-\beta_{t}{\color[rgb]{0.5703,0.1367,0.1367}b_{c}})$	Mahalanobis distance
$v_{c}(x,t)=\dot{\beta}_{t}{\color[rgb]{0.5703,0.1367,0.1367}b_{c}}+{\color[rgb]{0.5703,0.1367,0.1367}A_{c,t}}(x-\beta_{t}{\color[rgb]{0.5703,0.1367,0.1367}b_{c}})$	Gaussian Jeffreys divergence
$v_{\theta}(x,t)=\dot{\beta}_{t}{\color[rgb]{0.5703,0.1367,0.1367}b_{\theta}}+{\color[rgb]{0.5703,0.1367,0.1367}\bar{A}_{\theta,t}}(x-\beta_{t}{\color[rgb]{0.5703,0.1367,0.1367}b_{\theta}})$	Linear Fisher information

Table 1: Distance metrics and their corresponding velocity classes. We use

c

for a categorical condition and

\theta

for a continuous condition. All rows use the shared standard Gaussian source

\rho_{0}=\mathcal{N}(0,I)

except the starred geometry-template distance, which uses a source distribution concentrated around a low-dimensional manifold (Supplementary Section C ). The cosine and correlation rows equal

2r^{2}

times their conventional definitions. In the Mahalanobis row,

A_{t}

is shared across conditions. In the linear-Fisher row,

\bar{A}_{\theta,t}

is shared locally across nearby conditions (see Supplementary Section A.4 ).

Figure 3: Flow matching recovers distance metrics on a synthetic categorical dataset. ( A ) Pearson correlation between the unique off-diagonal entries of the estimated and ground-truth RDMs as a function of total sample size $N$ . Curves and error bars show means and sample standard deviations over five repetitions; the inset visualizes the five-condition dataset. ( B ) Ground-truth RDMs (top) and the corresponding flow-matching estimates using 2,000 total samples (bottom).

Figure 4: Flow matching performs best in estimating Fisher information on toy datasets. ( A ) Illustration of the Gaussian conditional-response dataset. ( B ) Representative linear Fisher-information estimates ( $d=300,N=1000$ ). ( C ) Mean relative absolute error (MRAE) versus total sample size at $d=300$ (left) and versus response dimension at $N=1{,}000$ (right). ( D ) Equal-weight Gaussian scale-mixture dataset. ( E ) Representative full Fisher-information estimates. ( F ) MRAE versus total sample size at $d=20$ (left) and versus response dimension at $N=7{,}000$ (right). Points and error bars in ( C ) and ( F ) show means and standard deviations over five repetitions.

Figure 5: Flow matching estimates Fisher information in mouse visual-cortex recordings. ( A ) Experimental paradigm: visual-cortex responses are recorded while mice view static gratings. ( B ) Session-wise mean held-out log likelihood relative to GKR. Gray lines connect estimates from the same session. Bars and error bars show the across-session mean and standard error of the mean (SEM). ( C ) Linear Fisher information from affine flow and full Fisher information from unconstrained flow versus grating orientation. Gray lines show six individual sessions, each from a different mouse; black curves and error bars show the across-session mean and SEM.

Figure 6: Flow matching best estimates time-resolved distances on OU-trial toy datasets. ( A ) Gaussian OU dataset with two conditions. ( B ) Representative time-resolved Gaussian Jeffreys-divergence estimates at $d=60$ and 260 trials per condition. ( C ) MRAE versus trials per condition at $d=60$ (left) and versus response dimension at 260 trials per condition (right). ( D ) Gaussian-mixture OU dataset. ( E ) Representative full Jeffreys-divergence estimates at $d=10$ and 260 trials per condition. ( F ) MRAE versus trials per condition at $d=10$ (left) and versus response dimension at 260 trials per condition (right). Points and error bars in ( C ) and ( F ) show means and standard deviations over five repetitions.

Figure 7: Flow matching estimates time-resolved distances from Allen Neuropixels recordings. ( A ) Four-second trial structure. ( B ) Mean test log-likelihood relative to the GKR method. Dots denote recording sessions, gray lines connect estimates from the same session, and bars show the mean across 32 sessions; triangles indicate two outliers above the displayed range. ( C ) Unconstrained-flow Jeffreys divergence averaged over the 28 unique pairs among the eight directions separately at each temporal frequency within each session (top), and 4-Hz direction RDMs averaged across sessions at the bin centers nearest $-0.5$ , $1.0$ , and $2.5$ seconds (bottom). Colors indicate temporal frequency. Curves and shaded bands show the across-session mean and SEM, respectively. ( D ) Time-to-time RDM correlations. Colors show mean Pearson correlations between RDMs at pairs of time points.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Experiment	Flow family	Width $\times$ depth	Batch	Learning rate
Categorical RDM	constrained and nonlinear	$128\times 3$	$3000$	$10^{-4}$ for correlation/cosine; $10^{-3}$ otherwise
Independent Gaussian Fisher	affine	$256\times 3$	$256$	$10^{-4}$
Independent mixture Fisher	affine and nonlinear	$256\times 3$	$256$	$10^{-4}$
Gaussian OU	affine	$256\times 3$	$256$	$10^{-4}$
Mixture OU	affine and nonlinear	$256\times 3$	$256$	$10^{-4}$
Stringer	affine and nonlinear	$256\times 3$	$256$	$10^{-4}$

Appendix

Table 2: Flow Matching training settings. Width

\times

depth gives the hidden width and number of hidden layers.

Benchmark	Mean-kernel widths	Covariance-kernel widths
Linear Fisher	$\lambda_{\mu}\in\{5,10,20\}$	$\lambda_{\Sigma}\in\{20{,}480,40{,}960,81{,}920\}$
Gaussian-mixture full Fisher	$\lambda_{\mu}\in\{2.25,9,36\}$	$\lambda_{\Sigma}\in\{144,576,2304\}$
Stringer recordings (initial grid)	$\lambda_{\mu}\in\{0.00125\cdot 2^{j}:j=0,\ldots,8\}$	$\lambda_{\Sigma}\in\{0.05\cdot 2^{j}:j=0,\ldots,7\}\cup\{\infty\}$
Gaussian OU	$\lambda_{\mu,t}\in\{1,4,16,64,256\}$	$\lambda_{\Sigma,t}\in\{256\cdot 4^{j}:j=0,\ldots,7\}\cup\{\infty\}$
Two-scale Gaussian-mixture OU	$\lambda_{\mu,t}\in\{1,4,16,64,256\}$	$\lambda_{\Sigma,t}\in\{16\cdot 4^{j}:j=0,\ldots,8\}\cup\{\infty\}$
Allen Neuropixels PCA-10	$\lambda_{\mu}\in\{0.015625\cdot 2^{j}:j=0,\ldots,6\}$ $\cup\{4,\infty\}$	$\lambda_{\Sigma}\in\{0.25,1,4,16,64\}$ $\cup\{1024,16384,262144,\infty\}$

Appendix

Table 3: Kernel-width grids used for the Wishart-process method. An infinite covariance width denotes a covariance that is constant over the continuous variable (e.g. orientation or trial time) within each condition.

Figure S1: Estimator accuracy across six representational distances on the five-condition categorical benchmark. Mean relative absolute RDM error is shown against total sample size for correlation, cosine, squared Euclidean, squared Mahalanobis, Gaussian Jeffreys divergence, and full Jeffreys divergence. Specialized estimators are compared with flow matching; the full-Jeffreys panel also includes telescoping density-ratio estimation (TRE) with 1, 8, or 16 bridges [ Rhodes et al., 2020 ] . Markers and error bars show the mean and one sample standard deviation across five repetitions.

Figure S2: Pearson correlation between the estimated and ground-truth curves. ( A ) Independent Gaussian responses and linear Fisher information. ( B ) Independent two-scale Gaussian-mixture responses and full Fisher information. ( C ) Gaussian Ornstein–Uhlenbeck trials and time-resolved Gaussian Jeffreys divergence. ( D ) Two-scale Gaussian-mixture Ornstein–Uhlenbeck trials and time-resolved full Jeffreys divergence. Method colors follow Figures 4 and 6 ; points and error bars show the mean and one sample standard deviation across five repetitions. Figure relates to Figures 4 and 6

Figure S3: Local conditional density fits across orientations in the representative Stringer session GT3. Rows show held-out responses in the first two principal components within circular $\pm 11.25^{\circ}$ windows centered at the displayed orientations. Columns compare held-out data, Bin + LW, GKR, the Wishart process, affine flow, and unconstrained flow. Crosses denote held-out observations; inner and outer ellipses or contours enclose 50% and 95% probability mass, respectively.

Figure S4: Conditional density and Fisher-information estimates for the Stringer recordings. ( A ) Held-out responses and fitted conditional densities in the first two principal components for an example session, colored by grating orientation. Gaussian estimators show fitted means and probability ellipses, whereas the unconstrained-flow panel shows generated samples. ( B ) Fisher-information estimates across the six sessions. GKR, Bin + LW, cross-fitted OLE, the Wishart process, and affine flow estimate linear Fisher information; unconstrained flow estimates full Fisher information. Thin gray curves show individual sessions, and colored curves and error bars show the pointwise mean and session-level standard error of the mean.

Figure S5: Adaptive Wishart kernel-width searches for the Stringer recordings. For each session, the validation surface shows Gaussian log likelihood relative to the session maximum over the mean width

\lambda_{\mu}

and covariance width

\lambda_{\Sigma}

; the white star marks the selected pair. The adjacent profile panel shows relative validation likelihood across mean widths (purple circles) and finite covariance widths (green triangles), with the constant-covariance endpoint shown by the orange cross.

Figure S6: Time-resolved conditional density trajectories for two drifting-grating conditions in the Allen Neuropixels recordings. Held-out responses and fitted distributions are shown in the first two principal components for the $(90^{\circ},2\,\mathrm{Hz})$ and $(180^{\circ},2\,\mathrm{Hz})$ conditions. Columns show held-out data, shared Bin + LW, Bin + LW, GKR, affine flow, and unconstrained flow. Marker shape identifies representative time bins from the beginning to the end of the analysis window, and color lightness indicates time from stimulus onset.

Figure S7: Frequency-specific time-resolved Gaussian Jeffreys divergence from classical estimators and affine flow across 32 Allen Neuropixels sessions. At each temporal frequency, Gaussian Jeffreys divergence is averaged over the 28 unique pairs among the eight directions within each session. Colors indicate the five temporal frequencies (

1

2

4

8

, and

15

Hz). The left panel shows Bin + LW (dashed curves) and GKR (dash-dotted curves); the right panel shows affine flow (solid curves). Curves show across-session means, shaded regions show the SEM, and solid and dashed vertical lines mark stimulus onset and offset. Bin + LW uses a 0.2-s time window.

Figure S8: Time-resolved all-condition RDMs across 32 Allen Neuropixels sessions. Each matrix is the elementwise mean across sessions for all 40 direction–frequency conditions at $-0.55$ , $0.95$ , or $2.45$ seconds relative to stimulus onset. Columns show Gaussian Jeffreys divergence from Bin + LW, GKR, and affine flow, followed by full Jeffreys divergence from unconstrained flow. Conditions are grouped by temporal frequency, with eight motion directions within each block.

Figure S9: Conditional density fits at four example times in an example Allen Neuropixels session. Held-out responses and fitted conditional densities are shown in the first two principal components for the $(90^{\circ},2\,\mathrm{Hz})$ and $(180^{\circ},2\,\mathrm{Hz})$ conditions. Rows correspond to windows centered at $-0.55$ , $0.45$ , $1.45$ , and $2.45$ seconds relative to stimulus onset; the displayed counts give the held-out observations for the two conditions. Columns compare held-out data, shared Bin + LW, Bin + LW, GKR, affine flow, and unconstrained flow.

Figure S10: Held-out likelihood across five Allen Neuropixels sessions after projection onto 10 principal components. Held-out log likelihood relative to GKR is shown for Bin + LW, GKR, the Wishart process, affine flow, and unconstrained flow. Gray lines connect estimates from the same session, black dots show individual sessions, and bars and error bars show the across-session mean and standard error of the mean.

Figure S11: Responses, held-out likelihood, and time-resolved representational geometry in the Allen two-photon recordings. ( A ) Four-second trial structure and baseline-corrected response heat maps for two drifting-grating conditions in a representative session. ( B ) Held-out log likelihood relative to GKR across 11 imaging sessions; gray lines connect estimates from the same session, and bars and error bars show the across-session mean and standard error of the mean. ( C ) Affine-flow Gaussian Jeffreys divergence averaged over condition pairs and 4-Hz direction RDMs at representative pre-stimulus, stimulus, and post-stimulus times. ( D ) Corresponding full Jeffreys divergence and RDMs from unconstrained flow. Curves and shaded regions summarize mean and SEM of the 11 sessions.

Figure S12: Similarity-constrained flow preserves a square template when estimating representational dissimilarity. Left: held-out observations from two noisy-square conditions, with solid curves showing their noiseless target boundaries. Middle and right: fits obtained using similarity-constrained and unconstrained flows, respectively, with the same side-length-one square source. In the fitted panels, the dashed black curve shows the original source boundary, and the solid colored curves show the condition-specific transported boundaries. Crosses denote held-out observations, and each panel title reports the Jeffreys divergence between the two condition distributions. The similarity-constrained fit preserves the square geometry, whereas the unconstrained fit distorts it.

Figure S13: Time-to-time RDM correlations across estimators in the Allen Neuropixels recordings. Panels compare Bin + LW, GKR, affine flow, and unconstrained flow. Each matrix shows Pearson correlations between RDMs at pairs of time points, computed within each recording session and then averaged across 32 sessions. Dashed lines mark stimulus onset at $0$ seconds and offset at $2$ seconds.

Neural Representation Geometry

Neural representation geometry studies how learned features are arranged and how that structure relates to model behavior and performance.

See all 256 papers

Feb 22, 2026

Partial Soft-Matching Distance for Neural Representational Comparison with Partial Unit Correspondence

Representational similarity metrics typically force all units to be matched, making them susceptible to noise and outliers common in neural representations. We extend the soft-matching distance to a partial optimal transport setting that allows some neurons to remain unmatched, yielding rotation-sensitive but robust correspondences. This partial soft-matching distance provides theoretical advantages -- relaxing strict mass conservation while maintaining interpretable transport costs -- and practical benefits through efficient neuron ranking in terms of cross-network alignment without costly iterative recomputation. In simulations, it preserves correct matches under outliers and reliably selects the correct model in noise-corrupted identification tasks. On fMRI data, it automatically excludes low-reliability voxels and produces voxel rankings by alignment quality that closely match computationally expensive brute-force approaches. It achieves higher alignment precision across homologous brain areas than standard soft-matching, which is forced to match all units regardless of quality. In deep networks, highly matched units exhibit similar maximally exciting images, while unmatched units show divergent patterns. This ability to partition by match quality enables focused analyses, e.g., testing whether networks have privileged axes even within their most aligned subpopulations. Overall, partial soft-matching provides a principled and practical method for representational comparison under partial correspondence.

Chaitanya Kapoor, Alex H. Williams, Meenakshi Khosla

Center for Neural Science, New York University, New York City, NY, 10003 · Center for Computational Neuroscience, Flatiron Institute, New York City, NY, 10003

Representational Similarity Analysis Neural Representation Geometry Latent Space Alignment

Ask PDF

Jun 8, 2026

A Theory on Flow Matching with Neural Networks

In this work, we develop theoretical foundation for flow matching with neural-network-parameterized conditional velocity fields. We establish convergence guarantees for gradient descent in the over-parameterized 2-layered ReLU neural network regime. We derive generalization bounds for the conditional velocity-field matching objective. Building on these results, we provide Wasserstein-distance guarantees for the samples generated by the induced flow. Our analysis is based on generalization bound for multi-task representation learning with unbounded losses, which may be of independent interest beyond flow-based generative modeling. These theoretical results are validated through extensive experiments on both synthetic and real-world image benchmarks.

Yihan He, Qishuo Yin, Yuan Cao +2

Princeton University, Princeton, NJ, USA. · The University of Hong Kong, Hong Kong. · Northwestern University, Evanston, IL, USA.

Neural Network Generalization Flow Matching Neural Network Optimization Flow-Based Generative Modeling Conditional Generative Modeling

Ask PDF

Mar 30, 2026

Geometry-aware similarity metrics for neural representations on Riemannian and statistical manifolds

Similarity measures are widely used to interpret the representational geometries used by neural networks to solve tasks. Yet, because existing methods compare the extrinsic geometry of representations in state space, rather than their intrinsic geometry, they may fail to capture subtle yet crucial distinctions between fundamentally different neural network solutions. Here, we introduce metric similarity analysis (MSA), a novel method which leverages tools from Riemannian geometry to compare the intrinsic geometry of neural representations under the manifold hypothesis. We show that MSA can be used to i) disentangle features of neural computations in deep networks with different learning regimes, ii) compare nonlinear dynamics, and iii) investigate diffusion models. Hence, we introduce a mathematically grounded and broadly applicable framework to understand the mechanisms behind neural computations by comparing their intrinsic geometries.

N Alex Cayco-Gajic, Arthur Pellegrino

Gatsby Unit, University College London

Neural Representation Geometry Representational Similarity Analysis Neural Network Interpretability

Ask PDF

Abstract

Neural Representation Geometry

Explore similar work

Partial Soft-Matching Distance for Neural Representational Comparison with Partial Unit Correspondence

A Theory on Flow Matching with Neural Networks

Geometry-aware similarity metrics for neural representations on Riemannian and statistical manifolds