Gaussian RBF Centered Kernel Alignment (CKA) in the Large Bandwidth Limit

Following the heuristic of [7], we use $\sigma^{\prime}=d_{X}\sigma$ instead of $\sigma$ in Eq. LABEL:eq:kernelOptions when computing the Gram matrix $K_{G(\sigma)}(X)$, where $d_{X}$ is the median distance between examples (rows) of $X$ (e.g., $\sigma=2$ describes a bandwidth equal to twice the median distance between examples). Equivalently, we divide all inter-example distances $|x_{i}-x_{j}|$ by $d_{X}$. This ensures invariance of Gaussian HSIC under isotropic scaling [7].

Our results rely crucially on mean-centering of features in Eqs. 2, 3. Indeed, the analog of Theorem 1 does not hold if, as in [3], features are not centered. This is easy to see by a direct calculation for the example in which feature matrix $X$ is the $2\times 2$ identity and feature matrix $Y$ is a $2\times 2$ of ones, except for a single off-diagonal $0$. Linear CKA equals $3/\sqrt{14}$ in that case, while Gaussian CKA is $1$ for all $\sigma$, as $X$ and $Y$ have identical distance matrices after scaling by the median distances $d_{X}$, $d_{Y}$ as described in the note in section 2.4. Details are provided in the Appendix.

Multiplying by $\sigma^{2}_{X}$, we obtain the connection in Eq. 8 between $\text{HSIC}(K_{G(\sigma)},L)$ and $\text{HSIC}(K_{E},L)$:

An identity analogous to Eq. 8 relates $\text{HSIC}(K_{G(\sigma)},K_{G(\sigma)})$ to its Euclidean version. In the latter case, the higher-order error terms in the HSIC expression are $O(\frac{1}{\sigma^{6}})$ before multiplying by $\sigma^{2}_{X}$, because they result from a product of a $O(\frac{1}{\sigma^{4}})$ error term from one of the two Gaussian Gram matrices with the $O(\frac{1}{\sigma^{2}})$ leading term in the other factor; multiplication by $\sigma^{4}_{X}$ is required to obtain a $O(\frac{1}{\sigma^{2}})$ residual, which will become a $\sigma^{2}_{X}$ factor after taking the square root in the CKA denominator (below).

Lemma 1 follows by dividing the HSIC components described above to form CKA as in Eq. 3:

The case of two Gaussian kernels with bandwidths approaching $\infty$ (whether equal to one another or not) follows by considering one kernel argument at a time, freezing the bandwidth of the other, and using the scalar triangle inequality. Details appear in the Appendix. $\Box$

Theorem 1 follows from Lemma 1 and Corollary 1.

We used sample OpenML data sets [22] (CC BY 4.0 license, https://creativecommons.org/licenses/by/4.0/) for classification (splice, tic-tac-toe, wdbc, optdigits, wine, dna) and regression (cpu, boston, Diabetes(scikit-learn), stock, balloon, cloud); for data sets with multiple versions, we used version 1. These data sets were selected based on two considerations only: smaller size, in order to eliminate the need for GPU acceleration and reduce environmental impact; and an absence of missing values, to simplify the development of internal feature representations for the CKA computations. Data set sizes are given in Table I.

Fully-connected neural networks (NN) with two hidden layers were used. Feature encodings $X$, $Y$ were the sets of activation vectors of the first and second hidden layers, respectively. Alternative NN widths, $w=16,32,64,128,256,512,1024$ were tested, with $w$ and $\frac{w}{4}$ hidden nodes in the first and second hidden layers. A configuration with three hidden layers, of sizes $128,32,8$, was also tested, for which $X$, $Y$ were the activation vectors from layers $1$ and $3$.

NN were trained using the MLPClassifier and MLPRegressor classes in scikit-learn [23], with cross-entropy or quadratic loss for classification and regression, respectively, ReLU activation functions, Glorot-He pseudorandom initialization [24], [25], Adam optimizer [26], learning rate of $0.001$, and a maximum of $2000$ training iterations. CKA was implemented in Python (https://docs.python.org/3/license.html), using NumPy [27] (https://numpy.org/doc/stable/license.html) and Matplotlib [28] (PSF license, https://docs.python.org/3/license.html). Experiments were performed on a workstation with an Intel i9-7920X (12 core) processor and 128GB RAM, under Ubuntu 18.04.5 LTS (GNU public license).

$50$ runs were performed for each pair $(D,w)$ of a data set $D$ and NN width $w$. In each run, a new NN model was trained on the full data set, starting from fresh pseudorandom initial parameter values. A training run was repeated if the resulting in-sample accuracy was below $0.8$ (classification) or if in-sample coefficient of determination was below $0.5$ (regression), but not if the optimization had not converged within the allowed number of iterations. After the network had been trained in a given run, $\text{CKA}_{\text{lin}}$ was computed once, and $\text{CKA}_{G(\sigma)}$ was computed for each bandwidth $\sigma=2^{p},\;p=-4,-3,\cdots,8$, for a total of $50$ $\text{CKA}_{\text{lin}}$ and $650$ $\text{CKA}_{G(\sigma)}$ evaluations per $(D,w)$ pair.

Compute time for the results presented in the paper was approximately $32$ hours, much of it on the $8\cdot 12\cdot(50+650)=67200$ CKA evaluations needed across the $8$ network widths (including the 128-32-8 three-hidden-layer configuration) and the $12$ data sets. Additional runs for validation required another $12$ hours. Several shorter preliminary runs were carried out for debugging and initial selection of the NN hyperparameters; configurations and threshold values were selected empirically in order to ensure a successful end to training after no more than a handful of runs in nearly all cases. Total compute time across all runs is estimated to have been $50$ hours.

CKA means and standard errors (SE = standard deviation divided by $\sqrt{50}$), and medians and standard error equivalent (SE = inter-quartile range divided by $\sqrt{50}$) of the ratio $\rho$ of Eq. 10 were computed across runs. We measured the magnitude of the discrepancy between $\text{CKA}_{G(\sigma)}$ and $\text{CKA}_{\text{lin}}$ by the base-$2$ logarithm of the relative difference between the two measures as in Eq. 11.

Theorem 1 implies that, for large $\sigma$, the logarithmic relative difference of Eq. 11 should decrease along a straight line of slope $-2$ as a function of $\log_{2}\sigma$, reflecting a $O(1/\sigma^{2})$ dependence. Accordingly, for each data set, we determined a $1/\sigma^{2}$ asymptote by extrapolating backward from the largest tested bandwidth value, $\sigma=2^{8}$. If $\sigma<2^{8}$, the predicted relative difference along the $1/\sigma^{2}$ asymptote is as in Eq. 12, where $r_{8}$ is the observed relative difference at $\sigma=2^{8}$.

We determined a $1/\sigma^{2}$ convergence onset bandwidth, $\sigma^{*}_{0}$, as the minimum bandwidth above which the observed $\log_{2}$ relative difference between Gaussian and linear CKA differs by less than $0.25$ from the predicted value along that data set’s $1/\sigma^{2}$ asymptote (Eq. 13).

Threshold values other than $0.25$ in Eq. 13 ($1,0.5,0.1$), corresponding to different tolerances for the log relative difference, yielded similar results in most cases in terms of the resulting relative $\sigma^{*}_{0}$ ranks of the different data sets.

Mean relative difference values between $\text{CKA}_{G(\sigma)}$ and $\text{CKA}_{\text{lin}}$ greater than $0.2$ ($\log_{2}$ values greater than $-2.3$), are observed for several data sets when $\sigma\leq 0.25$ (when $\log_{2}\sigma\leq-2$) in the case of classification; in fact, relative CKA difference values greater than $0.7$ ($\log_{2}\text{ rel.\ CKA diff.}>-0.5$, values in the Appendix) are observed for the dna data set when network width $w$ is $128$ or greater. This shows that noticeably nonlinear behavior of Gaussian CKA is quite possible for small bandwidth values. For some classification data sets, however (wdbc, wine), and most regression data sets, the relative CKA difference remains comparatively small for all bandwidths.

Fig. 1 shows confidence intervals for the means, of radius two standard errors, of the observed $\log_{2}$ relative difference between $\text{CKA}_{G(\sigma)}$ and $\text{CKA}_{\text{lin}}$ as a function of Gaussian bandwidth, $\sigma$ (Eq. 11). Convergence at the rate $1/\sigma^{2}$ for $\sigma\gg 1$ is observed (log-log slope of $-2$ in Fig. 1), as described in Theorem 1. Results are stable across the range of neural network configurations tested. Standard error of the relative CKA difference is observed to decrease as network width increases (e.g., Fig. 2), suggesting that wider networks are less sensitive to variations in initial parameter values.

The Diabetes(scikit-learn) data set stands out for having, by far, the smallest mean relative CKA difference among data sets tested, across much of the $\sigma$ range. The neural feature maps $X$ and $Y$ for that data set have nearly proportional linear Gram matrices $K_{\text{lin}}$ and $L_{\text{lin}}$, hence the linear CKA value is very close to $1$. Because the matrix of squared inter-example distances depends linearly on the linear Gram matrix, the Gaussian CKA value is also very close to $1$.

Our experimental results suggest that noticeably nonlinear behavior of Gaussian CKA occurs almost exclusively for bandwidths $\sigma<\rho$. Indeed, for all data sets tested, the observed mean relative difference between linear and Gaussian CKA is less than $0.01$ when $\sigma>\rho$, where $\rho$ is the ratio of maximum to median distance between feature vectors (Eq. 10). This confirms our finding in section 3.2 that $\text{CKA}_{G}(\sigma)$ differs little from $\text{CKA}_{\text{lin}}$ if $\sigma\gg\rho$. Furthermore, our results suggest that the threshold between nonlinear and linear regimes of Gaussian CKA is not substantially less than $\rho$: while mean relative CKA difference for $\sigma\geq 1$ peaked under $0.1$ for two-hidden-layer neural representations, across all data sets tested, relative difference values less than $0.1$ occur for the tic-tac-toe data set in the three-layer 128-32-8 representation only when $\sigma\geq 4$; the lower end of this bandwidth range is quite close to the median $\rho$ value of $4.2$ for that representation (which differs from the width-$64$ two-layer representation in Figs. 1, 2).

We find additional support for the view that $\rho$ approximates the boundary between nonlinear and linear regimes of Gaussian CKA in the fact that the convergence onset bandwidth, $\sigma^{*}_{0}$ (Eq. 13), is largest for the data sets of largest $\rho$: median $\log_{2}\sigma^{*}_{0}=5$, $4$, $4$ for balloon ($\rho\approx 16$), wdbc ($\rho\approx 11$), cpu ($\rho\approx 9$), respectively ($\sigma^{*}_{0}$ and $\rho$ values differ slightly among network widths and depths, but data sets of largest $\rho$ are the same). See Table III. Diabetes(scikit-learn) ties for third with $\log_{2}\sigma^{*}_{0}=4$; median $\log_{2}\sigma^{*}_{0}<4$ for all other data sets. Large differences in $\rho$ between data sets generally translate directly to differences between their $\sigma^{*}_{0}$ values. This is consistent with the analysis in section 3.2, as $\rho$ also controls terms of orders higher than $2$ in the exponential power series.

Our experimental results as a whole confirm $1/\sigma^{2}$ convergence of $\text{CKA}_{G(\sigma)}$ to $\text{CKA}_{\text{lin}}$ as $\sigma\rightarrow\infty$ (Theorem 1), and support the use of the representation eccentricity, $\rho$ (Eq. 10), as a useful heuristic upper bound on the range of bandwidths for which Gaussian CKA can display nonlinear behavior.