Neural networks learn to magnify areas near decision boundaries
In machine learning, there is a long history of trying to build neural
networks that can learn from fewer example data by baking in strong geometric
priors. However, it is not always clear a priori what geometric constraints are
appropriate for a given task. Here, we consider the possibility that one can
uncover useful geometric inductive biases by studying how training molds the
Riemannian geometry induced by unconstrained neural network feature maps. We
first show that at infinite width, neural networks with random parameters
induce highly symmetric metrics on input space. This symmetry is broken by
feature learning: networks trained to perform classification tasks learn to
magnify local areas along decision boundaries. This holds in deep networks
trained on high-dimensional image classification tasks, and even in
self-supervised representation learning. These results begins to elucidate how
training shapes the geometry induced by unconstrained neural network feature
maps, laying the groundwork for an understanding of this richly nonlinear form
of feature learning.
Jacob A. Zavatone-Veth, Sheng Yang, Julian A. Rubinfien, Cengiz Pehlevan
# Neural networks learn to magnify areas near decision boundaries

###### Abstract

We study how training molds the Riemannian geometry induced by neural network feature maps. At infinite width, neural networks with random parameters induce highly symmetric metrics on input space. Feature learning in networks trained to perform classification tasks magnifies local areas along decision boundaries. These changes are consistent with previously proposed geometric approaches for hand-tuning of kernel methods to improve generalization.

## 1 Introduction

In a series of influential papers, Amari and Wu proposed that one could improve the generalization performance of support vector machine (SVM) classifiers through data-dependent transformations of the kernel to expand the Riemannian volume element near decision boundaries (Amari & Wu, 1999; Williams et al., 2007; Wu & Amari, 2002). This proposal was based on the idea that this local magnification of areas improves class discriminability (Amari & Wu, 1999; Burges, 1999; Cho & Saul, 2011).

Over the past decade, SVMs have largely been eclipsed by neural networks, whose ability to flexibly learn features from data is believed to underlie their superior generalization performance (LeCun et al., 2015; Zhang et al., 2021). Previous works have explored some aspects of the geometry induced by neural networks feature maps with random parameters (Amari et al., 2019; Benfenati & Marta, 2023; Cho & Saul, 2009; 2011; Hauser & Ray, 2017; Poole et al., 2016; Zavatone-Veth & Pehlevan, 2022), but have not characterized data-dependent changes in representational geometry over training.

In this work, we explore the possibility that neural networks learn to enhance local input discriminability automatically over the course of training. Our primary contributions are:

* In SS4, we study general properties of the metric induced by shallow fully-connected neural networks. Next, in SS4.2, we compute the volume element and curvature of the metric induced by infinitely wide shallow networks with Gaussian weights and smooth activation functions, showing that it is spherically symmetric.
* In SS5, we empirically show that training shallow networks on simple classification tasks expands the volume element along decision boundaries, consistent with the hand-engineered modifications proposed by Amari and Wu. In SS6, we provide evidence that deep residual networks trained on more complex tasks behave similarly.

In total, our results provide a preliminary picture of how feature learning shapes local input discriminability.

## 2 Preliminaries

We begin by introducing the basic idea of the Riemannian geometry of feature space representations. Our setup and notation largely follow Burges (1999), which in turn follows the conventions of Dodson & Poston (1991).

### Feature embeddings as Riemannian manifolds

Consider \(d\)-dimensional data living in some submanifold \(\mathcal{D}\subseteq\mathbb{R}^{d}\). Let the _feature map_\(\mathbf{\Phi}:\mathbb{R}^{d}\rightarrow\mathcal{H}\) be a map from \(\mathbb{R}^{d}\) to some separable Hilbert space \(\mathcal{H}\) of possibly infinite dimension \(n\), with \(\mathbf{\Phi}(\mathcal{D})=\mathcal{M}\subseteq\mathcal{H}\). We index input space dimensions by Greek letters \(\mu,\nu,\rho,\ldots\in[d]\) and feature space dimensions by Latin letters \(i,j,k,\ldots\in[n]\). We use the Einstein summation convention; summation over all repeated indices is implied.

Assume that \(\mathbf{\Phi}\) is \(\mathcal{C}^{k}\) for \(k\geq 3\), and is everywhere of rank \(r=\min\{d,n\}\). If \(r=d\), then \(\mathcal{M}\) is a \(d\)-dimensional \(\mathcal{C}^{k}\) manifold immersed in \(\mathcal{H}\). If \(k=\infty\), then \(\mathcal{M}\) is a smooth manifold. In contrast, if \(r<d\), then \(\mathcal{M}\) is a \(d\)-dimensional \(\mathcal{C}^{k}\) manifold submersed in \(\mathcal{H}\). The flat metric on \(\mathcal{H}\) can then be pulled back to \(\mathcal{M}\), with components

\[g_{\mu\nu}=\partial_{\mu}\Phi_{i}\partial_{\nu}\Phi_{i}, \tag{1}\]
where we write \(\partial_{\mu}\equiv\partial/\partial x^{\mu}\). If \(r=d\) and the pullback metric \(g_{\mu\nu}\) is full rank, then \((\mathcal{M},g)\) is a \(d\)-dimensional Riemannian manifold (Burges, 1999; Dodson & Poston, 1991). However, if the pullback \(g_{\mu\nu}\) is a degenerate metric, as must be the case if \(r<d\), then \((\mathcal{M},g)\) is a singular semi-Riemannian manifold (Benfenati & Marta, 2023b; Kupeli, 2013). In this case, if we let \(\sim\) be the equivalence relation defined by identifying points with vanishing pseudodistance, the quotient \((\mathcal{M}/\sim,g)\) is a Riemannian manifold (Benfenati & Marta, 2023b). Unless noted otherwise, our results will focus on the non-singular case. We denote the matrix inverse of the metric tensor by \(g^{\mu\nu}\), and we raise and lower input space indices using the metric.

If we define the feature kernel \(k(\mathbf{x},\mathbf{y})=\Phi_{i}(\mathbf{x})\Phi_{i}(\mathbf{y})\) for \(\mathbf{x},\mathbf{y}\in\mathcal{D}\), then the resulting metric can be written in terms of the kernel as \(g_{\mu\nu}=(1/2)\partial_{x_{\mu}}\partial_{x_{\nu}}k(\mathbf{x},\mathbf{x})- [\partial_{y_{\mu}}\partial_{y_{\nu}}k(\mathbf{x},\mathbf{y})]_{\mathbf{y}= \mathbf{x}}\). This formula applies even if \(n=\infty\), giving the metric induced by the feature embedding associated to a suitable Mercer kernel (Burges, 1999).

### Volume element and curvature

With this setup, \((\mathcal{M},g)\) is a Riemannian manifold, hence we have at our disposal a powerful toolkit with which we may study its geometry. We will focus on two geometric properties of \((\mathcal{M},g)\). First, the volume element is given by

\[dV=\sqrt{\det g}\,d^{d}x, \tag{2}\]

where the factor \(\sqrt{\det g}\) measures how local areas in input space are magnified by the feature map (Amari & Wu, 1999; Burges, 1999; Dodson & Poston, 1991). Second, we consider the intrinsic curvature of the manifold, which is characterized by the Riemann tensor \(R^{\mu}_{\nu\alpha\beta}\)(Dodson & Poston, 1991). If \(R^{\mu}_{\nu\alpha\beta}=0\), then the manifold is intrinsically flat. As a tractable measure, we focus on the Ricci curvature scalar \(R=g^{\beta\nu}R^{\alpha}_{\nu\alpha\beta}\), which measures the deviation of the volume of an infinitesimal geodesic ball in the manifold from that in flat space (Dodson & Poston, 1991).

In the singular case, we can compute the volume element on \(\mathcal{M}/\sim\) at a given point by taking the square root of the product of the non-zero eigenvalues of the degenerate metric \(g_{\mu\nu}\) at that point (Benfenati & Marta, 2023b). However, the curvature in this case is generally not straightforward to compute; we will therefore leave this issue for future work.

### Shallow neural network feature maps

In this work, we consider a particular class of feature maps: those given by the hidden layer representations of neural networks (Benfenati & Marta, 2023b; Cho & Saul, 2009; 2011; Hauser & Ray, 2017; LeCun et al., 2015; Lee et al., 2018; Matthews et al., 2018; Neal, 1996; Williams, 1997). We will mostly focus on shallow fully-connected neural networks, i.e., those with only a single hidden layer followed by readout. Concretely, such a feature map is of the form

\[\Phi_{j}(\mathbf{x})=n^{-1/2}\phi(\mathbf{w}_{j}\cdot\mathbf{x}+b_{j}) \tag{3}\]

for weights \(\mathbf{w}_{j}\), biases \(b_{j}\), and an activation function \(\phi\). For convenience, we abbreviate the Euclidean inner product on feature or input space by -, e.g., \(\mathbf{w}\cdot\mathbf{x}=w_{\mu}x_{\mu}\). In this case, the feature space dimension \(n\) is equal to the number of hidden units, and is referred to as the _width_ of the hidden layer. In (3), we scale the components of the feature map by \(n^{-1/2}\) such that the associated kernel \(k(\mathbf{x},\mathbf{y})=\Phi_{i}(\mathbf{x})\Phi_{i}(\mathbf{y})\) and metric (4) have the form of averages over hidden units, and therefore should be well-behaved at large widths (Neal, 1996; Williams, 1997).

We will assume that \(\phi\) is \(\mathcal{C}^{k}\) for \(k\geq 3\), so that this feature map satisfies the smoothness conditions required in the setup above. We will also assume that the activation function and weight vectors are such that the Jacobian \(\partial_{\mu}\Phi_{j}\) is full-rank, i.e., is of rank \(\min\{d,n\}\). Then, the shallow network feature map satisfies the required conditions for the feature embedding to be a (possibly singular) Riemannian manifold. These conditions extend directly to deep fully-connected networks formed by composing feature maps of the form (3) (Benfenati & Marta, 2023b; Hauser & Ray, 2017).

## 3 Related works

Having established the geometric preliminaries of SS2, we can give a more complete overview of related works. As introduced above, our hypothesis for how the Riemannian geometry of neural network representations changes during training is directly inspired by the work of Amari & Wu (1999). In that and subsequent works (Amari & Wu, 1999; Williams et al., 2007; Wu & Amari, 2002), they proposed to modify the kernel of an SVM as \(k(\mathbf{x},\mathbf{y})=h(\mathbf{x})h(\mathbf{y})k(\mathbf{x},\mathbf{y})\) for some positive scalar function \(h(\mathbf{x})\) chosen such that the magnification factor \(\sqrt{\det g}\) is large near the SVM's decision boundary. Concretely, they proposed to fit an SVM with some base kernel \(k\), choose \(h(\mathbf{x})=\sum_{\mathbf{v}\in\text{SV}(k)}\exp[-\|\mathbf{x}-\mathbf{v}\|^ {2}/2\tau^{2}]\) for \(\tau\) a bandwidth parameter and \(\text{SV}(k)\) the set of support vectors for \(k\), and then fit an SVM with the modified kernel \(\tilde{k}\). Here, \(\|\cdot\|\) denotes the Euclidean norm. This process could then be iterated, yielding a sequence of modified kernels. They found that this method can improve generalization performance relative to the original kernel (Amari & Wu, 1999; Williams et al., 2007; Wu & Amari, 2002). This approach is a hand-designed form of iterative feature learning.

The geometry induced by common kernels was investigated in detail by Burges (1999), who established a broad range of technical results. He showed that translation-invariant kernels of the form \(k(\mathbf{x},\mathbf{y})=k(\|\mathbf{x}-\mathbf{y}\|^{2})\) yield flat, con
stant metrics,1 and gave a detailed characterization of polynomial kernels \(k(\mathbf{x},\mathbf{y})=(\mathbf{x}\cdot\mathbf{y})^{q}\). Cho and Saul (2011) subsequently analyzed the geometry induced by arc-cosine kernels, i.e., the feature kernels of infinitely-wide shallow neural networks with threshold-power law activation functions \(\phi(x)=\max\{0,x\}^{q}\) and random parameters (Cho and Saul, 2009). Our results on infinitely-wide networks for general smooth activation functions build on these works.

Footnote 1: It is interesting to note that this holds for the simplest form of a method for learning data-adaptive kernels recently proposed by Radhakrishnan et al. (2022); see Appendix F.

The representational geometry of deep networks with random Gaussian parameters in the limit of large width and depth was studied by Poole et al. (2016), and in later work by Amari et al. (2019). These works tie into a broader line of research on infinite-width limits of deep neural networks in which inference and prediction is captured by a kernel machine (Bordelon and Pehlevan, 2022; Daniely et al., 2016; Lee et al., 2018; Matthews et al., 2018; Neal, 1996; Williams, 1997; Yang, 2019; Yang and Hu, 2021; Zavatone-Veth and Pehlevan, 2022; Zavatone-Veth et al., 2021). Our results on the representational geometry of wide shallow networks with smooth activation functions build on these ideas, particularly those relating activation function derivatives to input discriminability (Daniely et al., 2016; Poole et al., 2016; Zavatone-Veth and Pehlevan, 2021, 2022).

Particularly closely related to our work are several recent papers that aim to study the curvature of neural network representations. Benfenati and Marta (2023b); Hauser and Ray (2017) discuss formal principles of Riemannian geometry in deep neural networks, but do not characterize how training shapes the geometry. Kaul and Lall (2020) aimed to study the curvature of metrics induced by the outputs of pretrained classifiers. However, their work is limited by the fact that they estimate input-space derivatives using inexact finite differences under the strong assumption that the input data is confined to a _known_ smooth submanifold of \(\mathbb{R}^{d}\). Recent works by Kuhnel et al. (2018); Shao et al. (2018); Wang and Ponce (2021) have studied the Riemannian geometry of the latent representations of deep generative models. Finally, in very recent work Benfenati and Marta (2023a) have used the geometry induced by the full input-output mapping to reconstruct iso-response curves of deep networks. In contrast, our work focuses on hidden representations, and seeks to characterize the representational manifolds themselves.

## 4 Representational geometry of shallow neural network feature maps

We begin by studying general properties of the Riemannian metrics induced by shallow neural networks feature maps.

### Finite-width networks

We first consider finite-width networks with fixed weights, assuming that \(n\geq d\). Writing \(z_{j}=\mathbf{w}_{j}\cdot\mathbf{x}+b_{j}\) for the preactivation of the \(j\)-th hidden unit, the general formula (1) for the metric yields

\[g_{\mu\nu}=\frac{1}{n}\phi^{\prime}(z_{j})^{2}w_{j\mu}w_{j\nu}. \tag{4}\]

This metric has the useful property that \(\partial_{\alpha}g_{\mu\nu}\) is symmetric under permutation of its indices, hence the formula for the Riemann tensor simplifies substantially (Appendix A).

Then, using the Leibniz formula for determinants, we show in Appendix B that the determinant of the metric can be expanded as a sum over \(d\)-tuples of hidden units:

\[\det g=\frac{1}{n^{d}d!}M_{j_{1}\cdots j_{d}}^{2}\phi^{\prime}(z_{j_{1}})^{2} \cdots\phi^{\prime}(z_{j_{d}})^{2}, \tag{5}\]

where

\[M_{j_{1}\cdots j_{d}}=\det\begin{pmatrix}w_{j_{1}1}&\cdots&w_{j_{1}d}\\ \vdots&\ddots&\vdots\\ w_{j_{d}1}&\cdots&w_{j_{d}d}\end{pmatrix} \tag{6}\]

is the minor of the weight matrix obtained by selecting units \(j_{1},\ldots,j_{d}\). For the error function \(\phi(x)=\operatorname{erf}(x/\sqrt{2})\), \(\det g\) expands as a superposition of Gaussian bump functions, one for each tuple of hidden units (B.55). This is reminiscent of Amari and Wu's approach, which yields a Gaussian contribution to \(\sqrt{\det g}\) from each support vector (SS3).

We can also derive similar expansions for the Riemann tensor and Ricci scalar. The resulting expressions are rather unwieldy, so we give their general forms only in Appendix B.3. However, in two dimensions the situation simplifies, as the Riemann tensor is completely determined by the Ricci scalar (Dodson and Poston, 1991; Misner et al., 2017) (Appendix B.1). In this case, we have the compact expression

\[(\det g)^{2}R=-\frac{3}{n^{3}}M_{jk}^{2}M_{ij}M_{ik}\\ \times\phi^{\prime}(z_{i})^{2}\phi^{\prime}(z_{j})\phi^{\prime}(z_ {k})\phi^{\prime\prime}(z_{j})\phi^{\prime\prime}(z_{k}). \tag{7}\]

This shows that in \(d=2\) the curvature acquires contributions from each triple of distinct hidden units, hence if \(n=2\) we have \(R=0\). This follows from the fact that the feature map is in this case a change of coordinates on the two-dimensional manifold (Dodson and Poston, 1991).

### Geometry of infinite shallow networks

We now characterize the metric induced by infinite-width networks (\(n\to\infty\)) with Gaussian weights and biases

\[\mathbf{w}_{j}\sim_{\text{i.i.d}}\mathcal{N}(\mathbf{0},\sigma^{2}\mathbf{I}_ {d});\quad b_{j}\sim_{\text{i.i.d.}}\mathcal{N}(0,\zeta^{2}), \tag{8}\]
as commonly chosen at initialization (LeCun et al., 2015; Lee et al., 2018; Matthews et al., 2018; Poole et al., 2016; Yang, 2019; Yang & Hu, 2021). For such networks, the hidden layer representation is described by the neural network Gaussian process (NNGP) kernel (Lee et al., 2018; Matthews et al., 2018; Neal, 1996; Williams, 1997):

\[k(\mathbf{x},\mathbf{y}) =\lim_{n\rightarrow\infty}n^{-1}\mathbf{\Phi}(\mathbf{x})\cdot \mathbf{\Phi}(\mathbf{y})\] \[=\mathbb{E}_{\mathbf{w},b}[\phi(\mathbf{w}\cdot\mathbf{x}+b)\phi (\mathbf{w}\cdot\mathbf{y}+b)]. \tag{9}\]

This kernel also completely describes the representation after training for networks in the lazy regime (Bordelon & Pehlevan, 2022; Yang & Hu, 2021).

In Appendix C, we show that the metric associated with the NNGP kernel, \(g_{\mu\nu}=\mathbb{E}_{\mathbf{w},b}[\phi^{\prime}(\mathbf{w}\cdot\mathbf{x} +b)^{2}w_{\mu}w_{\nu}]\), can be written more illuminatingly as

\[g_{\mu\nu}=e^{\Omega(\|\mathbf{x}\|^{2})}[\delta_{\mu\nu}+2\Omega^{\prime}(\| \mathbf{x}\|^{2})x_{\mu}x_{\nu}], \tag{10}\]

where the function \(\Omega(\|\mathbf{x}\|^{2})\) is defined via

\[e^{\Omega(\|\mathbf{x}\|^{2})}=\sigma^{2}\mathbb{E}_{z\sim\mathcal{N}(0, \sigma^{2}\|\mathbf{x}\|^{2}+\zeta^{2})}[\phi^{\prime}(z)^{2}]. \tag{11}\]

For these results, we must also assume that \(\phi\) and its (weak) derivatives satisfy suitable boundedness assumptions for \(\Omega\) to be twice-differentiable (Daniely et al., 2016). Therefore, like the metrics induced by other dot-product kernels, the NNGP metric has the form of a projection (Burges, 1999). Such metrics have determinant

\[\det g=e^{\Omega d}(1+2\|\mathbf{x}\|^{2}\Omega^{\prime}) \tag{12}\]

and Ricci scalar

\[R=- \frac{3(d-1)e^{-\Omega}(\Omega^{\prime})^{2}\|\mathbf{x}\|^{2}}{ (1+2\|\mathbf{x}\|^{2}\Omega^{\prime})^{2}}\] \[\times\left[d+2+2\|\mathbf{x}\|^{2}\left((d-2)\Omega^{\prime}+2 \frac{\Omega^{\prime\prime}}{\Omega^{\prime}}\right)\right]. \tag{13}\]

Thus, all geometric quantities are spherically symmetric, depending only on \(\|\mathbf{x}\|^{2}\). Thanks to the assumption of independent Gaussian weights, the geometric quantities associated to the shallow Neural Tangent Kernel and to the deep NNGP will share this spherical symmetry (Appendix D) (Lee et al., 2018; Matthews et al., 2018; Yang, 2019; Yang & Hu, 2021). This generalizes the results of Cho & Saul (2011) for threshold-power law functions to arbitrary smooth activation functions. The relation between Gaussian norms of \(\phi^{\prime}\) and input discriminability indicated by this result is consistent with previous studies (Daniely et al., 2016; Poole et al., 2016; Zavatone-Veth & Pehlevan, 2021, 2022). In short, unless the task depends only on the input norm, the geometry of infinite-width networks will not be linked to the task structure.

### Examples

In Appendix C.2, we evaluate the geometric quantities of the NNGP for certain analytically tractable activation functions. The resulting expressions for \(\sqrt{\det g}\) and \(R\) are rather lengthy, so we discuss only their qualitative behavior here. For the error function \(\phi(x)=\mathrm{erf}(x/\sqrt{2})\), \(R\) is negative for all \(d>1\), and both \(R\) and \(\sqrt{\det g}\) are monotonically decreasing functions of \(\|\mathbf{x}\|\) for all \(\zeta\) and \(\sigma\). For monomials \(\phi(x)\propto x^{q}\) for integer \(q>1\), \(\sqrt{\det g}\) is a monotonically increasing function of \(\|\mathbf{x}\|^{2}\), while \(R\) is again non-positive. However, in this case the behavior of \(R\) depends on whether or not bias terms are present: if \(\zeta=0\), then \(R\) is a non-decreasing function of \(\|\mathbf{x}\|^{2}\) that diverges towards \(-\infty\) as \(\|\mathbf{x}\|^{2}\downarrow 0\), while if \(\zeta>0\), \(R\) may be non-monotonic in \(\|\mathbf{x}\|^{2}\). In Figure 1, we illustrate this behavior, and show convergence of the empirical geometry of finite networks with random Gaussian parameters to the infinite-width results.

Figure 1: Convergence of geometric quantities for finite-width networks with Gaussian random parameters to the infinite-width limit. **a**. The magnification factor \(\sqrt{\det g}\) (_left_) and Ricci scalar \(R\) (_right_) as functions of the input norm \(\|\mathbf{x}\|\) for networks with \(\phi(x)=\mathrm{erf}(x/\sqrt{2})\). Empirical results for finite networks, computed using (5) and (B.15) are shown in blue, with solid lines showing the mean and shaded patches the standard deviation over \(25\) realizations of random Gaussian parameters. In all cases, \(\sigma=\zeta=1\). The infinite-width result is shown as a black dashed line. **b**. As in **a**, but for normalized quadratic activation functions \(\phi(x)=x^{2}/\sqrt{3}\).
## 5 Changes in shallow network geometry during training

We now consider how the geometry of the pullback metric changes during training. Changes in the volume element and curvature during gradient descent training are challenging to study analytically, because models for which the learning dynamics are solvable--deep linear networks (Saxe et al., 2013)--trivially yield flat, constant metrics. More generally, the dependence of the metric on the instantaneous configuration of parameters makes it difficult to gain intuition for its evolution over training, even for two-dimensional inputs.

### Wide Bayesian neural networks

We can make slightly more analytical progress for Bayesian neural networks at large but finite width. This setting is convenient because there is a fixed parameter posterior; one does not need to solve kernel or metric dynamics through time (Bordelon and Pehlevan, 2022). In Appendix E, we use recent results on perturbative feature-learning corrections to the NNGP kernel (Roberts et al., 2022; Zavatone-Veth et al., 2021) to compute corrections to the posterior mean of the volume element. In general, it is not possible to evaluate these corrections in closed form (Zavatone-Veth et al., 2021). For networks with monomial activation functions, no bias terms, and linear readout constrained to interpolate a single training example \((\mathbf{x}_{a},\mathbf{y}_{a})\), we can show that the correction to \(\sqrt{\det g}\) is maximal for \(\mathbf{x}\parallel\mathbf{x}_{a}\), and minimal for \(\mathbf{x}\perp\mathbf{x}_{a}\) (E.38). The sign of the correction is positive or negative depending on whether the second moment of the prior predictive is greater than or less than the norm of the output, respectively. For example, if we train on a single point from the XOR task, \((1,1)\mapsto 0\), \(\sqrt{\det g}\) will be contracted maximally along \(x_{1}=x_{2}\). This simple case

Figure 2: Evolution of the volume element over training in a network trained to classify points separated by a sinusoidal boundary \(y=\frac{3}{5}\sin(7x-1)\) (single hidden layer with 5 hidden units (top), 20 hidden units (mid), and 250 hidden units (bottom)). Red lines indicate the decision boundaries of the network. See Appendix G.1 for experimental details. More hidden units offer better approximation to the sinusoid curve.
shows how interpolating a single point shapes the network's global representational geometry.

### Changes in representational geometry for networks trained on two-dimensional toy tasks

Thus, given the intractability of studying changes in geometry analytically, we resort to numerical experiments. For details of our numerical methods, see Appendix G.

To build intuition, we first consider networks trained on simple two-dimensional tasks, for which we can directly visualize the input space. We first consider a simple two-dimensional binary classification task with sinusoidal boundary, inspired by that considered in the original work of Amari & Wu (1999). We train networks with sigmoidal activation functions of varying widths to perform this task, and visualize the resulting geometry over the course of training in Figure 2. At initialization, the peaks in the volume element lack a clear relation to the structure of the task, with approximate rotational symmetry at large widths as we would expect from SS4.2. As the network's decision boundary is gradually molded to conform to the true boundary, the volume element develops peaks in the same vicinity. At all widths, the final volume elements are largest near the peaks of the sinusoidal decision boundary. At small widths, the shape of the sinusoidal curve is not well-resolved, but at large widths there is a clear peak in the close neighborhood of the decision boundary. This result is consistent with the proposal of Amari & Wu (1999).

In Appendix G, Figure G.5, we plot the Ricci curvature for these trained networks. Even for these small networks, the curvature computation is computationally expensive and numerically challenging. Over training, it evolves dynamically, with task-adapted structure visible at the end of training. However, the patterns here are harder to interpret than those in the volume element.

### Changes in representational geometry for shallow networks trained to classify MNIST digits

We now provide evidence that a similar phenomenon is present in networks trained to classify MNIST images. We give details of these networks in Appendix G.2; note that all reach above 95% train and test accuracy within 200 epochs.

In Figure 3, we plot the induced volume element at synthetic images generated by linearly interpolating between two input images (see Appendix G for details). We emphasize that linear interpolation in pixel space of course does not respect the structure of the image data, and results in unrealistic images. However, this approach has the advantage of being straightforward, and also illustrates how small Euclidean perturbations are expanded by the feature map (Novak et al., 2018). At initialization, the volume element varies without clear structure along the interpolated path. However, as training progresses, areas near the center of the path, which roughly aligns with the decision boundary, are expanded, while those near the endpoints defined by true training examples remain relatively small. This is again consistent with the proposal of Amari & Wu (1999). We provide additional visualizations of this behavior in Appendix G.2.

To gain an understanding of the structure of the volume element beyond one-dimensional slices, in Figure 3 we also plot its value in the plane spanned by three randomly-selected example images, at points interpolated linearly within their convex hull. Here, we only show the end of training; in Appendix G.2 we show how the volume element in this plane changes over the course of training. The edges of the resulting ternary plot are one-dimensional slices like those shown in the top row of Figure 3, and we observe consistent expansion of the volume element along these paths. The volume element becomes large near the centroid of the triangle, where multiple decision boundaries intersect.

Because of the computational complexity of estimating the curvature--the Riemann tensor has \(d^{2}(d^{2}-1)/12\) independent components (Dodson & Poston, 1991; Misner et al., 2017)--and its numerical sensitivity (Appendix G.1), we do not attempt to estimate it for this high-dimensional task.

## 6 Extensions to deep networks

Thus far, we have focused on the geometry of the feature maps of single-hidden-layer neural networks. However, these analyses can also be applied to deeper networks, regarding the representation at each hidden layer as defining a feature map, and study how the geometry changes with depth (Benfenati & Marta, 2023b; Hauser & Ray, 2017). As a simple version of this, in Figure G.6 we consider a network with three fully-connected hidden layers trained on the sinusoid task. The metrics induced by the feature maps of all three hidden layers all show the same qualitative behavior as we observed in the shallow case in Figure 2: areas near the decision boundary are magnified. As one progresses deeper into the network, the contrast between regions of low and high magnification factor increases.

As a more realistic example, we consider deep residual networks (ResNets) (He et al., 2016) trained to classify the CIFAR-10 image dataset (Krizhevsky, 2009). To make the feature map differentiable, we replace the rectified linear unit (ReLU) activation functions used in standard ResNets with Gaussian error linear units (GELUs) (Hendrycks & Gimpel, 2016). With this modification, we achieve comparable test accuracy (92%) with a ResNet-34--the largest model we can consider given computational constraints--to that obtained with ReLUs (Appendix G.3). Importantly, the feature map defined by the input-to-final-hidden-layer map
ping of a ResNet-34 gives a submersion of CIFAR-10, as the input images have \(32\times 32\times 3=3072\) pixels, while the final hidden layer has 512 units. Empirically, we find that the Jacobian of this mapping is full-rank (Appendix G.3); we therefore consider the volume element on \((\mathcal{M}/\sim,g)\) defined by the product of the non-zero eigenvalues of the degenerate pullback metric (SS2, Appendix G.3).

In Figures 4, we visualize the resulting geometry in the same way we did for networks trained on MNIST, along 1-D interpolated slices and in a 2-D interpolated plane (see Appendix G.3 for details and additional figures). In the 1-D slices, we see a clear trend of large volume elements near decision boundaries, as we observed for shallow networks. However, in two dimensions the picture is less clear. The decision boundaries are more complicated than for MNIST, reflecting the more complex structure of the task. This also highlights the deficiency of our approach of linear interpolation in pixel space, which we further discuss and illustrate in Appendix G.3. We observe some magnification of areas in the vicinity of decision boundaries, though here it is harder to interpret all forms of structure that are present. Thus, even in this more realistic setting, we observe shaping of geometry over training that appears consistent with the proposal of Amari and Wu (1999).

## 7 Discussion

To conclude, we have shown that training on simple tasks shapes the Riemannian geometry induced by neural network representations by magnifying areas along decision boundaries, consistent with the proposal of Amari and Wu for geometrically-inspired kernel learning (Amari and Wu, 1999; Williams et al., 2007; Wu and Amari, 2002). Our results on the geometry induced by the NNGP kernel provide a preliminary picture of the geometric priors of neural networks, and our experimental results begin to show how representational geometry is shaped over the course of training. These results are relevant to the broad goal of leveraging non-Euclidean geometry in deep learning (Bronstein et al., 2021; Weber et al., 2020). We now discuss the limitations of our work,

Figure 3: _Top panel_: \(\log_{10}(\sqrt{\det g})\) induced at interpolated images between 7 and 6 by a single-hidden-layer fully-connected network trained to classify MNIST digits. _Bottom panel_: Digit class predictions and \(\log_{10}(\sqrt{\det g})\) for the plane spanned by MNIST digits 7, 6, and 1 at the final training epoch (200). Sample images are visualized at the endpoints and midpoint for each set. Each line is colored by its prediction at the interpolated region and end points. As training progresses, the volume elements bulge in the middle (near the decision boundary) and taper off when travelling towards endpoints. See Appendix G.2 for experimental details and Figure G.7 for images interpolated between other digits.
as well as directions for future study.

Perhaps the most important limitation of our work is the fact that we focus either on toy tasks with two-dimensional input domains, or on low-dimensional slices through high-dimensional domains. This is a fundamental limitation of how we have attempted to visualize the geometry. We are also restricted by computational constraints (see in particular Appendix G.3); to characterize the geometry of state-of-the-art network architectures, more efficient and numerically stable algorithms for computing these quantities must be developed.

An important question that we leave open for future work is whether expanding areas near decision boundaries generically improves generalization in deep neural networks, consistent with Amari & Wu (1999)'s original motivations. Indeed, it is easy to imagine a scenario in which the geometry is overfit, and the trained network becomes too sensitive to small changes in the input. This possibility is consistent with prior work on the sensitivity of deep networks Novak et al. (2018), and with the related phenomenon of adversarial vulnerability Goodfellow et al. (2014); Szegedy et al. (2013). Investigating these links will be an important objective for future work.

Because of the smoothness conditions required by the definition of the pullback metric and the requirement that \((\mathcal{M},g)\) be a differentiable manifold Benfenati & Marta (2023); Hauser & Ray (2017), the approach pursued in this work does not apply directly to networks with ReLU activation functions, which are not differentiable. Deep ReLU networks are continuous piecewise-linear maps, with many distinct activation regions Hanin & Rolnick (2019);b). Within each region, the corresponding linear feature map will induce a flat metric on the input space, but the magnification factor will vary from region to region. It will be interesting to investigate the resulting geometry in future work.

One possible area of application of our results is to the general problem of how to analyze and compare neural network representations Kornblith et al. (2019); Williams et al. (2021). Importantly, one could compute and plot the volume element induced by a feature map even when one

Figure 4: _Top panel_: \(\log_{10}(\sqrt{\det g})\) induced at interpolated images between a horse and a frog by ResNet34 trained to classify CIFAR-10 digits. _Bottom panel_: Digits classification of a horse, a frog, and a car. The volume element is the largest at the intersection of several binary decision boundaries, and smallest within each of the decision region. The one-dimensional slices along the edges of each ternary plot are consistent with the top panel. See Appendix G.3 for experimental details, Figure G.12 for linear interpolation and plane spanned by other classes, and how the plane evolves during training.
does not have access to explicit class labels. This could allow one to study networks trained with self-supervised learning, or even differentiable approximations to biological neural networks (Wang and Ponce, 2022). Exploring the rich geometry induced by these networks is an exciting avenue for future investigation.

## Acknowledgements

We thank Alexander Atanasov and Blake Bordelon for helpful comments on our manuscript. JAZV, CP and this research were supported by a Google Faculty Research Award and NSF DMS-2134157. The computations in this paper were run on the FASRC Cannon cluster supported by the FAS Division of Science Research Computing Group at Harvard University.
