Estimating Neural Representation Alignment from Sparsely Sampled Inputs and Features — part 2/5
Cleaned-text code points [15658, 31556). Bibliography omitted here; source.md includes it.
We find that the representation geometry affects the bias in the widely-used stimulus-corrected estimator CKA ^ S \widehat{\text{CKA}}_{S} . Assuming the variance of the ℋ ^ S \widehat{\mathcal{H}}_{S} estimates are negligible, and the activation variance is normalized neuron-wise, the bias of CKA ^ S \widehat{\text{CKA}}_{S} can be approximated as 𝔼 [ CKA ^ S ( k ( a ) , k ( b ) ) ] − CKA ( k ( a ) , k ( b ) ) ≈ [ 1 ( 1 + γ a − 1 Q a ) ( 1 + γ b − 1 Q b ) − 1 ] CKA ( k ( a ) , k ( b ) ) \mathbb{E}\left[\widehat{\text{CKA}}_{S}(k^{(a)},k^{(b)})\right]-\text{CKA}(k^{(a)},k^{(b)})\approx\\ \left[\frac{1}{\sqrt{\left(1+\frac{\gamma_{a}-1}{Q_{a}}\right)\left(1+\frac{\gamma_{b}-1}{Q_{b}}\right)}}-1\right]\text{CKA}(k^{(a)},k^{(b)}) (8) where γ a \gamma_{a} and γ b \gamma_{b} are the intrinsic dimensionalities of the two underlying representations, quantified by the participation ratio of the eigenvalues { λ i } \{\lambda_{i}\} of 𝐊 ¯ \bar{\mathbf{K}} in the infinite data limit: γ = ( ∑ i λ i ) 2 ∑ i λ i 2 . \gamma=\frac{\left(\sum_{i}\lambda_{i}\right)^{2}}{\sum_{i}\lambda_{i}^{2}}. From Equation 8 , we make the following observations:
• CKA ^ S \widehat{\text{CKA}}_{S} always underestimates the true CKA on average.
• The scale by which CKA ^ S \widehat{\text{CKA}}_{S} underestimates is independent of the alignment but only dependent on the intrinsic dimensionalities and neuron sample sizes.
• Having larger intrinsic dimensionalities contributes to the greater underestimation of CKA and requires more neuron samples to mitigate the bias.
Therefore, if intrinsic dimensionality is much greater than the number of feature samples, γ ≫ Q \gamma\gg Q , CKA ^ S \widehat{\text{CKA}}_{S} can take an arbitrarily small value even when two representations are perfectly aligned. On the other hand, our estimator is not affected by this issue. For more details, including the bias of CKA ^ 0 \widehat{\text{CKA}}_{0} , see Section B.2 .
## Stimulus-Neuron-Corrected Estimator
Unlike the previous two estimators, an unbiased estimator correcting for both finite stimulus and neuron sampling must be a function of populations ϕ ( a ) {\bm{\phi}}^{(a)} rather than their kernels.
In this paper, we develop a novel estimator ℋ ^ C \widehat{\mathcal{H}}_{C} that corrects for finite neuron sampling effects when two populations have correlated, e.g., identical neurons. We leave its derivation to Appendix A and report the final estimator here ℋ ^ C ( ϕ ( a ) , ϕ ( b ) ) = 1 P 3 ( P − 3 ) Q a ( Q b − δ a b ) × ∑ i j l m [ c i j i j − 2 P P − 2 ( c i j j l − c i i i j P − c i j j j P + c i i i i 2 P ) + P 2 ( P − 1 ) ( P − 2 ) ( c i j l m − c i i j l P − c j l i i P + c i i j j P 2 ) ] \widehat{\mathcal{H}}_{C}({\bm{\phi}}^{(a)},{\bm{\phi}}^{(b)})=\frac{1}{P^{3}\left(P-3\right)Q_{a}\left(Q_{b}-\delta_{ab}\right)}\\ \times\sum_{ijlm}\left[c_{ijij}-\frac{2P}{P-2}\left(c_{ijjl}-\frac{c_{iiij}}{P}-\frac{c_{ijjj}}{P}+\frac{c_{iiii}}{2P}\right)\right.\\ \left.+\frac{P^{2}}{\left(P-1\right)\left(P-2\right)}\left(c_{ijlm}-\frac{c_{iijl}}{P}-\frac{c_{jlii}}{P}+\frac{c_{iijj}}{P^{2}}\right)\right] (9)
where we introduced the tensor c i j l m c_{ijlm} defined as c i j l m ≔ ∑ α β 𝚽 i α ( a ) 𝚽 j α ( a ) \displaystyle c_{ijlm}\coloneqq\sum_{\alpha\beta}{\bm{\Phi}}_{i\alpha}^{(a)}{\bm{\Phi}}_{j\alpha}^{(a)} 𝚽 l β ( b ) 𝚽 m β ( b ) − δ a b ∑ α 𝚽 i α ( a ) 𝚽 j α ( a ) 𝚽 l α ( b ) 𝚽 m α ( b ) . \displaystyle{\bm{\Phi}}_{l\beta}^{(b)}{\bm{\Phi}}_{m\beta}^{(b)}-\delta_{ab}\sum_{\alpha}{\bm{\Phi}}_{i\alpha}^{(a)}{\bm{\Phi}}_{j\alpha}^{(a)}{\bm{\Phi}}_{l\alpha}^{(b)}{\bm{\Phi}}_{m\alpha}^{(b)}. (10)
The second term in c i j l m c_{ijlm} encodes the finite neuron correction only when two populations have identical neurons ( a = b a=b ). Finally, we define our unbiased CKA estimator as CKA ^ C ( ϕ ( a ) , ϕ ( b ) ) = ℋ ^ C ( ϕ ( a ) , ϕ ( b ) ) ℋ ^ C ( ϕ ( a ) , ϕ ( a ) ) ℋ ^ C ( ϕ ( b ) , ϕ ( b ) ) . \displaystyle\widehat{\text{CKA}}_{C}\left({\bm{\phi}}^{(a)},{\bm{\phi}}^{(b)}\right)=\frac{\widehat{\mathcal{H}}_{C}({\bm{\phi}}^{(a)},{\bm{\phi}}^{(b)})}{\sqrt{\widehat{\mathcal{H}}_{C}({\bm{\phi}}^{(a)},{\bm{\phi}}^{(a)})\widehat{\mathcal{H}}_{C}({\bm{\phi}}^{(b)},{\bm{\phi}}^{(b)})}}. (11) Several remarks are in order: • Our estimator ℋ ^ C \widehat{\mathcal{H}}_{C} reduces to the stimulus-corrected estimator ℋ ^ S \widehat{\mathcal{H}}_{S} of Song et al. ((2012)) when two populations ϕ ( a ) {\bm{\phi}}^{(a)} and ϕ ( b ) {\bm{\phi}}^{(b)} have independent neurons, generalizing previous results.
• When two distinct populations compared, our CKA estimator CKA ^ C \widehat{\text{CKA}}_{C} corrects the bias in the denominator of CKA ^ S \widehat{\text{CKA}}_{S} and their numerators remain identical. However, in certain cases, e.g. trial-to-trial similarity between two measurements of a single population, the numerator of CKA ^ S \widehat{\text{CKA}}_{S} also becomes biased.
• While the estimator looks complicated, it can be efficiently implemented in terms of einsum operations. An implementation is provided in Appendix D .
• Note that while the ℋ \mathcal{H} -estimator we derive is unbiased, the proposed CKA-estimator is still biased due to non-linear operations (multiplication, square-root) involving ℋ ^ C \widehat{\mathcal{H}}_{C} . Nevertheless, our CKA estimator is the least biased among the available estimators. Full analysis on this matter can be found in Section B.3 . Also, as shown in the rest of the paper, we empirically find that the effect of this bias is much smaller compared to the effect of bias due to finite neuron sampling. One can also mitigate this bias by empirically averaging ℋ ^ \widehat{\mathcal{H}} over multiple trials and using that to estimate CKA. If N N estimates of ℋ ^ \widehat{\mathcal{H}} are available from independent trials, then this bias falls like 𝒪 ( 1 / N ) \mathcal{O}\left(1/N\right) ( Section B.3 ). We will denote the CKA estimator derived from the empirical average of N N number of ℋ ^ \widehat{\mathcal{H}} by CKA N ^ \widehat{\text{CKA}^{N}} .
## Linear example
Next, we numerically test all three estimators on a simple synthetic dataset with a known CKA value. We consider d − d- dimensional stimuli which are drawn from the distribution 𝐱 ∼ 𝒩 ( 0 , 𝐈 d ) {\mathbf{x}}\sim\mathcal{N}(0,{\mathbf{I}}_{d}) . We define a linear population of the form ϕ ( a ) ( 𝐱 ) = 𝐱 ⊤ 𝐰 ( a ) {\bm{\phi}}^{(a)}({\mathbf{x}})={\mathbf{x}}^{\top}{\mathbf{w}}^{(a)} , where each weight is drawn from the distribution 𝐰 ( a ) ∼ 𝒩 ( 0 , 𝚺 a ) {\mathbf{w}}^{(a)}\sim\mathcal{N}(0,{\bm{\Sigma}}_{a}) and corresponds to a single neuron. Similarly, we define the population ϕ ( b ) ( 𝐱 ) = 𝐱 ⊤ 𝐰 ( b ) {\bm{\phi}}^{(b)}({\mathbf{x}})={\mathbf{x}}^{\top}{\mathbf{w}}^{(b)} where 𝐰 ( b ) ∼ 𝒩 ( 0 , 𝚺 b ) {\mathbf{w}}^{(b)}\sim\mathcal{N}(0,{\bm{\Sigma}}_{b}) . Since the entire stimuli and population distributions are known, the true CKA in Equation 4 can be evaluated exactly: CKA = tr ( 𝚺 a 𝚺 b ) tr 𝚺 a 2 tr 𝚺 b 2 . \displaystyle\text{CKA}=\frac{\tr(\bSigma_a \bSigma_b)}{\sqrt{\tr{\bm{\Sigma}}_{a}^{2}\tr{\bm{\Sigma}}_{b}^{2}}}. (12) We consider 𝚺 a = 𝚺 b = 𝐈 d {\bm{\Sigma}}_{a}={\bm{\Sigma}}_{b}={\mathbf{I}}_{d} in which case both populations perfectly align with a CKA = 1 \text{CKA}=1 . As indicated by Equation 8 , however, despite the perfect alignment of populations, the stimulus-corrected estimator CKA ^ S \widehat{\text{CKA}}_{S} can be arbitrarily small depending on the number of neurons sampled Q Q , and the dimensionality γ \gamma , which is exactly d d in this example. We numerically confirm these findings in Figure 2 and find that both estimators become highly sensitive when a limited number of neurons are observed. On the other hand, our estimator is able to recover the true CKA even at small neuron samples.
In this example, we observe that CKA N ^ \widehat{\text{CKA}^{N}} ’s and CKA ^ \widehat{\text{CKA}} ’s almost completely overlap in all three estimators ( Figure 2 lighter vs. darker lines), except when Q Q is very small, CKA N ^ C \widehat{\text{CKA}^{N}}_{C} and CKA ^ C \widehat{\text{CKA}}_{C} diverge a bit. This means that the bias correction for ℋ \mathcal{H} contributes more to the bias correction of CKA than correcting the bias from the non-linear operations on the ℋ \mathcal{H} estimates. The next section shows that this pattern is observed in neural data as well.
## Practical applications in neuroscience
In most of the brain recordings, the observed neurons are samples of much larger population. Here, we use electrophysiological data of the three key cortical regions for visual processing, V1, V4, and IT, in the order of processing cascade. It has been shown that the neurons in V1 are simple filters, but the later regions V4 and IT are sensitive to semantic information ( Hubel & Wiesel, (1962) ; Majaj et al., (2015) ; Hung et al., (2005) ; Cadena et al., (2024) ) . Papale et al. ((2025)) presented 20,000 natural images of 1,854 object categories from the THINGS dataset ( Hebart et al., (2019) ) to two monkeys and recorded neural responses with multiple electrode arrays over V1, V4, and IT, totaling 1046 electrodes in one monkey and 960 in the other. The measurement value is the average spiking voltage levels of neurons adjacent to a given electrode averaged over a small time window immediately following an image presentation. We view the electrodes and images as rows and columns, respectively.
### Brain-to-Brain alignment
Here use CKA to benchmark the similarity of a given brain region across individual animals (see Figure 1 a). Such brain-to-brain comparisons provide an essential reference for evaluating model-to-brain alignments. In our analysis, we first compare two disjoint sets of electrodes from the same brain region within a single animal. The ground truth in this within-region comparison is a perfect alignment (CKA = 1) on average. However, as shown in Figure 3 a, both the naive estimator CKA ^ 0 \widehat{\text{CKA}}_{0} and the stimulus-corrected estimator CKA ^ S \widehat{\text{CKA}}_{S} significantly underestimate the true similarity, with their estimates strongly dependent on the number of neurons sampled. In contrast, our proposed estimator CKA ^ C \widehat{\text{CKA}}_{C} closely approximates the true value even with very limited neuron sampling, thereby highlighting the risk that conventional CKA estimates may fail to detect even perfect alignment when undersampling is present.
We further extend the analysis by measuring the similarity of each brain region across animals. As depicted in Figure 3 b, CKA ^ C \widehat{\text{CKA}}_{C} yields reliable estimates over a range of neuron sampling sizes, whereas the conventional estimators remain highly sensitive to the sample size. Our results show that V1 representations are relatively conserved across individuals (with CKA around 0.6), IT representations exhibit lower similarity (around 0.3), and V4 lies in between (approximately 0.45). The similarity decreases over the layers of visual processing.
### Brain-to-Model alignment
The alignment between brain representations and artificial neural network layers is a topic of considerable interest. It has been repeatedly observed in the literature that the early visual regions, e.g., V1, have representations that are strikingly similar to the early layers of CNNs, and the later regions, e.g., IT, are aligned with the deeper layers ( Yamins et al., (2014) ; Schrimpf et al., (2018) ; Nonaka et al., (2021) ) . As depicted in Figure 1 b, here we use the CKA estimators to compare the representations between brain regions (V1, V4, and IT) and layers of convolutional neural networks (CNNs), namely ResNet18 and CorNet-s ( He et al., (2016) ; Kubilius et al., (2018) ) . We then see how these CKA estimates compare against the brain-to-brain CKA ^ C \widehat{\text{CKA}}_{C} from the previous section.
We make the same observation of the V1-early layer alignment and IT-late layer alignment from all of the estimators. In the absolute scale, the V1-early layer alignment is generally higher, but the IT-late layer alignment is generally lower ( Figure 4 left three plots). However, CKA ^ C \widehat{\text{CKA}}_{C} estimates that the alignment between animals also decreases by a similar factor. This indicates that in the relative scale, the model-to-brain alignments for all three regions are similarly close to the animal-to-animal alignments, an insight that cannot be reliably found with the other estimators.
We then test the sensitivity of the estimators to the number of electrodes. In Figure 4 (right three plots), we pick the best alignment layer for each brain region and see how this alignment value changes with the number of electrodes. We observe that all estimators are similar for comparison with V1, meaning that there is only a small bias in CKA ^ 0 \widehat{\text{CKA}}_{0} and CKA ^ S \widehat{\text{CKA}}_{S} . However, for the comparison with V4 and IT, the gap between the estimators is generally large, meaning we need our estimator CKA ^ C \widehat{\text{CKA}}_{C} for a reliable measurement. The CKA ^ 0 \widehat{\text{CKA}}_{0} and CKA ^ S \widehat{\text{CKA}}_{S} values are not converged even when all available electrodes are used.
### Object disentanglement
Beyond inter-region comparisons, our estimator enables a novel application of CKA: quantifying the disentanglement of semantic information across object categories in the brain. Similar to the representation disentanglement observed in deeper layers of CNNs, the IT cortical region is believed to encode high-level semantic features that distinguish between objects ( Yamins et al., (2014) ) . In this analysis, we estimate the CKA between pairs of object categories within a single brain region, where the columns of the measurement matrix correspond to stimulus images and the rows correspond to the shared neurons (as illustrated in Figure 1 b). A lower CKA in this context suggests greater semantic disentanglement.
We observe that the estimated CKA between object categories generally decreases over V1, V4, and IT, indicating a gradual semantic disentanglement over these regions ( Figure 5 ). Interestingly, natural object pairs get separated faster than artificial object pairs. This pattern is more strongly observed in CKA ^ C \widehat{\text{CKA}}_{C} than in CKA ^ S \widehat{\text{CKA}}_{S} . Interestingly, there are many cases where CKA ^ S \widehat{\text{CKA}}_{S} indicates high disentanglement (small CKA) but ours CKA ^ C \widehat{\text{CKA}}_{C} indicates low disentanglement (large CKA). This is observed in V1 for all pairwise comparisons: Our estimator consistently estimates CKA value near 1, suggesting V1 does not encode semantic information, whereas the heavily biased CKA ^ S \widehat{\text{CKA}}_{S} spuriously suggests otherwise. Also, in IT, CKA ^ C \widehat{\text{CKA}}_{C} is near 1 for artificial objects pairs, but the bias in CKA ^ S \widehat{\text{CKA}}_{S} spuriously suggests the objects are disentangled.
## Trial-to-Trial Similarity
As our final application, we evaluate the trial-to-trial similarity of neural recordings from the same brain region on the same stimuli set as depicted in Figure 1 c. We use the neural recordings from Papale et al. ((2025)) of visual cortical areas V1, V2, and IT over 30 30 trials on 100 100 images. This case differs from the previous analyses since a single population ϕ ( a ) {\bm{\phi}}^{(a)} is compared across two trials ( ϕ 1 ( a ) {\bm{\phi}}^{(a)}_{1} and ϕ 2 ( a ) {\bm{\phi}}^{(a)}_{2} ). In this case, CKA ^ S \widehat{\text{CKA}}_{S} becomes biased in its numerator as well as its denominator, which our estimator accounts for correctly.