implicit bias, stochastic gradient descent, mean-field limit, ReLU neural networks, feature learning.

1 Introduction↩︎

Overparameterized neural networks generally admit infinitely many interpolating solutions for a given training data set. Yet optimization algorithms such as stochastic gradient descent (SGD) consistently select solutions that generalize remarkably well. Understanding the principles governing this selection mechanism is a central question in modern learning theory. One approach to this question is through the notion of implicit bias, which posits that optimization algorithms not only minimize the training error but also systematically favor certain minimizers over others. For example, in linear models, gradient-based optimization is known to select minimum-norm solutions. More generally, a large body of work has characterized the implicit bias of gradient descent and SGD through variational principles involving margin maximization, reproducing kernel Hilbert space norms, total variation of curvatures, and other notions of complexity [1]–[4]. Despite this substantial progress, our understanding of the implicit bias of nonlinear neural networks remains incomplete, particularly in feature-learning regimes.

A natural framework for studying implicit bias in feature-learning dynamics is the mean-field regime [5]–[8]. In this regime, a wide two-layer network is represented by a probability measure over neurons, and the training dynamics are described by a Wasserstein gradient flow on the space of probability measures. Moreover, when stochastic noise is incorporated into the dynamics, the resulting evolution admits a unique stationary measure that characterizes the long-time behavior of noisy SGD. This raises the following fundamental question: What is the structure of the stationary measure and of the solution function represented by this stationary measure?

1.0.0.1 Main Contributions

We study multivariate regression with wide two-layer ReLU networks trained by noisy stochastic gradient descent with weight decay in the mean-field regime. Our main results, illustrated in Figure 1, are as follows:

• In Theorem 2, we show that, as the noise level tends to zero, the solution selected by noisy SGD converges, irrespective of the strength of weight decay, to a continuous piecewise affine function. More precisely, this function is affine on each connected component of the complement of a finite hyperplane arrangement.

• We further show that the number of hyperplanes is bounded above by \(2\mathcal{P}-1\), where \(\mathcal{P}\) denotes the number of linear dichotomies realizable on the training inputs. Thus, the linear-region complexity of the learned solution is controlled by the combinatorial geometry of the training data set.

• Next, in Theorem 3, we show that the input weights and biases of the trained network collapse onto at most \(2\mathcal{P}- 1\) distinct directions. Consequently, the trained network has effective width at most \(2\mathcal{P}- 1\).

• Finally, in Theorem 4, we establish a strong non-redundancy property of these effective neurons. Specifically, distinct effective neurons induce distinct ternary activation patterns on the training inputs.

Our work builds on the analysis of [9], who characterized the solution functions selected by noisy SGD in the mean-field regime for univariate ReLU networks. We extend this analysis to the multivariate setting and, in addition, characterize the structure of the stationary measure, thereby describing the network parameters at convergence. In Section 2 we place our results in the context of recent advances on mean-field dynamics, noisy SGD, implicit bias, and sparse or finite-support solutions. To our knowledge, our results provide the first derivation of effective width collapse as an implicit bias of noisy SGD for multivariate regression with general training data, together with a problem-dependent characterization of the learned solution.

1.0.0.2 Organization

The remainder of the article is organized as follows. Section 2 gives an overview of related work. Section 3 provides the necessary preliminaries, introducing the mean-field analysis of two-layer networks and discussing the implementation of ReLU activation functions within this framework. Section 4 presents our main results, and Section 5 discusses their interpretation, technical aspects, and limitations. Section 6 outlines the proof strategy and key components. Section 7 presents numerical experiments illustrating the main results. Supporting results and technical proofs are collected in the appendices.

2 Related Work↩︎

2.0.0.1 Implicit Bias for Wide Neural Networks

The implicit bias of gradient-based parameter optimization in wide neural network depends significantly on the training regime, in particular, on the scale of network parametrization and initialization relative to the network width [10], [11]. A significant line of work has studied the implicit bias in the lazy or kernel regime [12]–[15], where the network parameters remain close to their initialization throughout training and no feature learning takes place. In this regime, the implicit bias is characterized as minimization of a reproducing kernel Hilbert space norm [1], whose explicit form has been investigated in series of works [16]–[20].

A distinct and important training regime is the mean-field regime, where the network parameters are modeled as an interacting particle system whose evolution is described by a nonlinear PDE over probability measures [5]–[8]. This setting is of particular interest because it captures many nontrivial dynamical phenomena, including feature learning. However, compared to the kernel regime, much less is known about the convergence of training and the resulting implicit bias. The challenges in analyzing the mean-field dynamics are discussed by [21]. Nevertheless, several important advances have been made in characterizing the implicit bias in the mean-field regime. In a classification setting, [2] showed that if the network parameters converge in direction, the converged direction is a maximizer for an \(\mathcal{F}_1\) norm maximum margin problem. [22] provided an analysis for hinge loss and linearly separable data with a spherically symmetric input distribution. Assuming infinitely many training samples, they showed that the ReLU network implements a linear classifier throughout training.

The work most closely related to ours is that of [9], whose analysis provides the main foundation for the present work. They studied noisy SGD in a setting where the global convergence of the corresponding mean-field PDE to a stationary measure had previously been established by [5]. For univariate ReLU networks, they showed that, in the vanishing-noise limit, the stationary measure corresponds to a piecewise linear function with at most three knots between any two consecutive training inputs. A detailed comparison between their results and ours is provided in Section 5.

2.0.0.2 Explaining Sparsity from Training Dynamics

The phenomenon that training dynamics lead to sparse features has been widely observed empirically [10], [23], [24]. One attempt to explain this phenomenon is through early alignment [23], whereby, under infinitesimal initialization, the network parameters align along a small number of directions during the initial phase of training. Motivated by this phenomenon, a line of work has investigated the full training dynamics and established effective width collapse [25]–[31]. Building on the connection between gradient flow and maximum-margin optimization in homogeneous networks for classification [3], [4], [32], [33] demonstrated an implicit bias toward low-rank input weights. Moreover, [34] showed that, for univariate classification problems whose labels are generated by a teacher network of width \(r\), the learned function has effective width \(O(r)\). For Gaussian cluster data and linear classification loss, [35] showed that, with infinitely many training samples, the trained ReLU network yields a linear decision boundary.

Besides [9], existing works establishing effective width collapse as a consequence of the training dynamics rely on restrictive assumptions on the training data. In classification, these include orthogonally separable [25], [27], symmetric separable [26], nearly orthogonal [32], [33], and positively correlated [30] training data. In multivariate regression, [28] assume that the training inputs are mutually orthogonal, [29] consider labels generated by a single neuron, and [31] analyze a training set consisting of three data points. In addition, these works typically require the parameter initialization to have sufficiently small scale or to satisfy a balancedness condition. By contrast, we establish effective width collapse under minimal assumptions on the training data and parameter initialization. The tradeoff is that our analysis relies on the network being sufficiently wide, allowing us to exploit the mean-field framework.

2.0.0.3 Explaining Sparsity from Optimal Solutions

Our results complement prior work on sparse or finite-support solutions of static optimization problems that capture the learning objective. The setting closest to ours is that of shallow ReLU networks with unbounded width and Euclidean norm regularization of the network parameters. By the homogeneity of the ReLU activation, this is equivalent to regularizing the \(\mathcal{F}_1\) functional norm [36]. For univariate ReLU networks with a free skip connection, [37] showed that, the minimum-\(\mathcal{F}_1\)-norm interpolator is unique and can be expressed as the sum of an affine function and a ReLU network of width at most \(M-1\), where \(M\) is the number of training samples. For multivariate data, [38], [39] showed that every minimum-\(\mathcal{F}_1\)-norm interpolator admits a representation as a ReLU network with effective width bounded by a constant \(K\) that depends on the input dimension and sample size. By contrast, we derive an effective-width bound for the solution selected by noisy SGD. Our bound is strictly sharper than those of [38], [39]; see a detailed comparison in Section 5.

A different regularizer, which penalizes only the Euclidean or \(\ell_p\) norm of the network weights but not the biases, has been studied both for unbounded-width networks [40]–[43] and finite-width networks [44], [45]. In both settings, it has been shown that there exist optimal solutions that can be represented as ReLU networks with width at most \(K\), where \(K\) is a data-dependent constant.

In classification, another widely studied optimization problem is the maximum-margin problem. This line of work is motivated by the observation of [3], [4] that, under suitable conditions, training homogeneous neural networks converges in direction to a KKT point of a maximum-margin problem. Subsequent work has investigated the structure of the maximum-margin solutions and their connection to sparse feature representations [46]–[50].

While these works provide a detailed understanding of the solutions of the associated optimization problem, it remains largely open under what conditions a training algorithm converges to such solutions. Our results bridge this gap by showing that effective width collapse arises not only as a property of static variational problems, but also as a consequence of noisy SGD training dynamics.

3 Preliminaries↩︎

In this section, we recall preliminary results from the mean-field theory of wide neural networks and introduce a bounded, smoothed approximation of the ReLU neuron, which is needed to apply these mean-field results.

3.1 Mean-Field Analysis for Wide Two-Layer Networks↩︎

Consider fully connected networks with \(d\) input units, \(N\) hidden units, and one output unit, defined by \[\label{eq:model} h_\sigma(x,\Theta)= \frac{1}{N}\sum_{i=1}^N\sigma(x,\theta_i) ,\tag{1}\] where \(\sigma\colon \mathbb{R}^d \times \mathbb{R}^D\to \mathbb{R}\) is the unit (or neuron) activation, \(\theta_i\in \mathbb{R}^D\) denotes the parameter vector of the \(i\)th hidden unit, and \(\Theta=(\theta_i)_{i=1}^N \in \mathbb{R}^{ND}\) denotes the parameter of the entire network. To recover a ReLU network, we set \(\sigma(x,\theta)=a (w^\top x+b)_+\), where \(\theta=(a,w,b)\in \mathbb{R}\times \mathbb{R}^{d}\times\mathbb{R}= \mathbb{R}^{d+2}\) consists of the output weight \(a\), the input weights \(w\), and the bias \(b\), and where \((u)_+=\max(u,0)\). Assume that the parameters of the hidden units are initialized as i.i.d.samples from a fixed measure \(\rho_0\in \mathcal{P}_2^a(\mathbb{R}^D)\). Here and throughout, \(\mathcal{P}_2^a(\mathbb{R}^D)\) denotes the space of absolutely continuous probability measures on \(\mathbb{R}^D\) with finite second moments.

Consider a finite training data set \(\left\{(x_j,y_j)\right\}_{j=1}^M\subset \mathbb{R}^d \times \mathbb{R}\). We consider the mean squared error with \(\ell_2\)-regularization: \[L(\Theta)=\frac{1}{M}\sum_{j=1}^M (h_\sigma(x_j,\Theta)-y_j)^2 + \frac{\lambda}{N} \|\Theta\|^2.\] To minimize \(L\), we train the neural network using noisy stochastic gradient descent (noisy SGD), with parameter updates of the form \[\label{eq:SGD} \theta_{i}^{k+1} = \theta_i^k - 2 s_{k,N} \big(h_\sigma(\tilde{x}_k,\Theta^k)-\tilde{y}_k\big) \nabla_{\theta_i} \sigma(\tilde{x}_k, \theta_i^k) - 2s_{k,N} \lambda \theta_i^k + \sqrt{\frac{2s_{k,N}}{\beta}}g_i^k,\tag{2}\] where \(s_{k,N}\) is the step size, \((\tilde{x}_k,\tilde{y}_k)_{k\in {\mathbb{N}_0}}\) are i.i.d.samples drawn from the empirical distribution \(\frac{1}{M}\sum_{j=1}^M\delta_{(x_j,y_j)}\) induced by the training data, and \((g_i^k)_{i\in \left\{1,\ldots,N\right\},k\in \mathbb{N}}\) are i.i.d.Gaussian noise drawn from \(N(0,I_D)\). The parameter \(\beta^{-1}\) controls the noise level and is commonly referred to as the temperature. In this work, we focus on the vanishing-temperature regime, \(\beta^{-1} \to 0\), in which case 2 recovers the standard SGD algorithm [51].

The network model 1 depends on the parameter \(\Theta\) only through the empirical measure \(\tfrac{1}{N}\sum_{i=1}^N\delta_{\theta_i}\). In particular, the ordering of the hidden units does not affect the network output. This motivates a reformulation of the model in terms of probability measures \(\rho\in \mathcal{P}(\mathbb{R}^D)\): \[h_\sigma(x,\rho) = \int \sigma(x,\theta) \, d\rho.\] When \(\rho = \tfrac{1}{N}\sum_{i=1}^N \delta_{\theta_i}\), this representation recovers the finite-width network 1 .

Let \(\hat{\rho}_k^{(N)} = \tfrac{1}{N} \sum_{i=1}^N \delta_{\theta_i^k}\) denote the empirical distribution of the network parameters after \(k\) steps of noisy SGD 2 . Consider a step size scheme \(s_{k,N} = \varepsilon \xi(k\varepsilon)\), where \(\varepsilon=\varepsilon_N>0\) depends on the width \(N\) and \(\xi\colon [0,\infty) \to [0,\infty)\) is a sufficiently regular function, e.g., \(\varepsilon_N=N^{-1}\) and \(\xi\equiv 1\), so \(s_{k,N}= N^{-1}\). It is known that, for any fixed training time \(t\geq0\), the empirical parameter measure \(\hat{\rho}_{\lfloor t/\varepsilon \rfloor}^{(N)}\) converges weakly to \(\rho_t\) as \(\varepsilon \to 0\) and \(N\to\infty\), where \(\rho_t\) is the solution of the following partial differential equation with initial condition \(\rho_0\) (the parameter initialization distribution): \[\label{eq:pde} \begin{align} \partial_t \rho_t &= 2\xi(t)\nabla_\theta \cdot \Big(\rho_t \nabla_\theta \Psi^\lambda_\sigma(\theta;\rho_t) \Big) + 2\xi(t) \beta^{-1} \Delta_\theta\rho_t\\ \Psi^\lambda_\sigma(\theta;\rho) &= \frac{1}{M}\sum_{j=1}^M\Big(h_\sigma(x_j,\rho)-y_j\Big) \sigma(x_j,\theta) +\frac{\lambda}{2}\|\theta\|^2. \end{align}\tag{3}\] Here \(\nabla_\theta \, \cdot\) denotes the divergence operator and \(\Delta_\theta\) the Laplacian. The mean-field potential \(\Psi^\lambda_\sigma(\theta;\rho)\) is interpreted as the loss experienced by an individual neuron when all other neurons are represented collectively by the distribution \(\rho\). Thus, \(\hat{\rho}_{\lfloor t/\varepsilon \rfloor}^{(N)}\) may be viewed as a discrete-time, finite-particle approximation of \(\rho_t\). Note that in the step size scheme \(s_{k,N}=\varepsilon \xi(k\varepsilon)\), \(\varepsilon\) specifies the time discretization scale and \(\xi\) controls the speed of evolution of the PDE 3 .

Let \(\mathcal{H}\) denote the set of admissible densities of probability measures in \(\mathcal{P}_2^a(\mathbb{R}^D)\): \[\mathcal{H}= \left\{\rho \colon \mathbb{R}^{D}\to [0,+\infty) \colon \rho \text{ is measurable}, \int \rho(\theta)d\theta =1, M(\rho)<+\infty\right\},\] where \(M(\rho)=\int \|\theta\|^2\rho(\theta)d\theta\) is the second moment. For notational simplicity, we use \(\rho\) to denote both a measure and its density when the latter exists. Then PDE 3 can be interpreted as a Wasserstein gradient flow on \(\mathcal{H}\), minimizing the following free energy: \[\label{eq:energy} F_\sigma^{\beta,\lambda}(\rho) = \frac{1}{2M}\sum_{j=1}^M\Big(h_\sigma(x_j,\rho)-y_j\Big)^2 + \frac{\lambda}{2}M(\rho)- \frac{1}{\beta}H(\rho) ,\tag{4}\] where \(H(\rho)=-\int\rho(\theta)\log(\rho(\theta))d\theta\) is the entropy. Under suitable regularity assumptions on the activation function \(\sigma\) and the initialization measure \(\rho_0\), it has been shown that the free energy 4 admits a unique minimizer \(\rho^\ast\) in \(\mathcal{H}\), and that PDE 3 converges weakly to this minimizer. Moreover, the minimizer satisfies the following condition: \[\label{eq:minimizer} \rho^\ast(\theta) =Z_{\beta,\lambda,\rho^\ast}^{-1} \exp(-\beta \Psi^\lambda_\sigma(\theta;\rho^\ast)),\tag{5}\] where \(Z_{\beta,\lambda, \rho}=\int \exp(-\beta \Psi^\lambda_\sigma(\theta;\rho)) d\theta\) is the partition function. This condition follows from the Euler-Lagrange equation characterizing minimizers of the free energy \(F^{\beta,\lambda}_\sigma\) [5]. Notably, since \(\rho^\ast\) appears on both sides of 5 , the equation constitutes a self-consistency condition for the stationary measure. As the right-hand side has the form of a Boltzmann distribution, 5 is referred to as a Boltzmann fixed point condition.

The aforementioned results are summarized in the following theorem.

Thus, the function \(h_\sigma(\cdot,\rho^\ast)\) characterizes the implicit bias of noisy SGD training, and a fundamental goal of the present work is to understand its structure. This is challenging because the stationary measure \(\rho^\ast\) is known only implicitly through the Boltzmann fixed point condition 5 , a self-consistency equation that does not explicitly characterize either \(\rho^\ast\) or the function it represents.

3.2 Bounded and Smoothed Approximation of ReLU Neurons↩︎

Theorem 1 requires that the neuron map \(\sigma(x,\theta)\) is bounded and smooth in the parameter \(\theta\), assumptions that are not satisfied by the ReLU neuron \[\label{eq:relu-neuron} \sigma(x,\theta)=a(w^\top x+b)_+.\tag{6}\] In this subsection, we introduce a bounded and smooth approximation of the ReLU neuron that enables the application of Theorem 1. Our construction follows the approximation introduced by [9].

For \(m,\tau>0\), let \((\cdot)^{m,\tau}_+\) denote the \(C^4\) smooth, bounded approximation of the univariate ReLU activation, defined as follows: \[(u)_+^{m,\tau} = \begin{cases} S_\tau(u)=\log(1+e^{\tau u})/\tau,& u < S_\tau^{-1}(m^2)\\ \phi_{m,\tau}(u), & u\ge S_\tau^{-1}(m^2), \end{cases}\quad u\in \mathbb{R}.\] Thus, \((u)_+^{m,\tau}\) agrees with the softplus \(S_\tau(u)=\tau^{-1}\log(1+e^{\tau u})\) until it reaches the value \(m^2\), after which it is smoothly capped by the tail \(\phi_{m,\tau}\). The tail is chosen such that \(\|(\cdot)_+^{m,\tau}\|_\infty\leq 2m^2\) and \((\cdot)_+^{m,\tau}\in C^4(\mathbb{R})\) with bounded derivatives up to the fourth order.

We also introduce a smooth, bounded approximation of the identity map in order to make the neuron output 6 bounded in output weights \(a\). For \(m,\tau>0\) and \(n\in \mathbb{N}\), let \[v^{m,\tau} = \begin{cases} v,& \|v\|< m-\tau^{-1}\\ \psi_{m,\tau}(v), &\|v\|\ge m-\tau^{-1}, \end{cases}\quad v\in \mathbb{R}^n.\] The tail \(\psi_{m,\tau}\) is chosen so that \(\|(\cdot)^{m,\tau}\| \leq m\) and \((\cdot)^{m,\tau}\in C^4(\mathbb{R}^n,\mathbb{R}^n)\) with bounded derivatives up to the fourth order. In addition, we assume that the tail functions \(\phi_{m,\tau}\) and \(\psi_{m,\tau}\) satisfy suitable regularity properties; in particular, for every \(m>0\) the limits \(\lim_{\tau\to+\infty}(u)_+^{m,\tau}\) and \(\lim_{\tau\to+\infty}v^{m,\tau}\) exist. For brevity, we defer the detailed construction to Appendix 9.

Define the \(m\)-truncated, \(\tau\)-smoothed ReLU neuron with \(m,\tau>0\): \[\sigma^{m,\tau}(x,\theta)= a^{m,\tau} ( x^\top w^{m,\tau}+b )_+^{m,\tau},\quad x\in \mathbb{R}^d,\theta=(a,w,b)\in \mathbb{R}^{d+2}.\] The truncation of the input weights \(w\) is introduced only for technical reasons. By construction, the approximated ReLU neuron converges pointwise to the ReLU neuron as \(m,\tau\to+\infty\), i.e., \(\lim_{m,\tau\to+\infty} \sigma^{m,\tau}(x,\theta)=a(w^\top x+b)_+\) for every \((x,\theta)\). Moreover, \(\sigma^{m,\tau}\) satisfies the assumptions of Theorem 1. Let \(\rho_{\beta,\lambda,m,\tau}^\ast\) denote the stationary measure of the mean-field PDE 3 associated with the neuron \(\sigma^{m,\tau}\), regularization parameter \(\lambda\), and temperature \(\beta^{-1}\). With a slight abuse of notation, let \[h^\ast_{\beta,\lambda,m,\tau}(\cdot)=h_{\sigma^{m,\tau}}(\cdot, \rho_{\beta,\lambda,m,\tau}^\ast)\] denote the corresponding network function.

The goal of this work is to study the stationary measure \(\rho^\ast_{\beta,\lambda,m,\tau}\) and the corresponding solution function \(h^\ast_{\beta,\lambda,m,\tau}\) in order to characterize the implicit bias of noisy SGD in both parameter space and function space. We impose the following assumptions on \((\beta,\lambda,m,\tau)\).

The next result removes the smoothing parameter \(\tau\) by showing that both the stationary measure and the corresponding solution function have pointwise limits as \(\tau\to+\infty\). Its proof is inspired by [9] and is deferred to Appendix 10.1.

We are interested in the limits of \(\rho^\ast_{\beta,\lambda,m}\) and \(h^\ast_{\beta,\lambda,m}\) as \(m,\beta\to+\infty\), where \(m\to+\infty\) recovers the ReLU activation and \(\beta\to+\infty\) corresponds to the vanishing-temperature limit of noisy SGD. The existence of these limits \[\label{eq:mblimit} \lim_{m,\beta\to+\infty}\rho^\ast_{\beta,\lambda,m}(\cdot),\quad \lim_{m,\beta \to+\infty}h^\ast_{\beta,\lambda,m}(\cdot)\tag{7}\] remains open even in the univariate setting [9]. Several obstacles arise. As \(\beta\to+\infty\), the free energy 4 generally loses strict convexity and may therefore have multiple minimizers. Moreover, as \(m\to+\infty\), the activation function becomes unbounded, causing the free energy functional to potentially lose lower semicontinuity. Consequently, the standard arguments used to establish the existence of minimizers no longer apply [52]. These difficulties are intrinsic to the free energy and the unboundedness of the ReLU activation, rather than artifacts of the particular approximation introduced above.

Motivated by these obstacles, we do not assume that the limits in 7 exist. Instead, we analyze all accumulation points arising in the joint limit \(\beta,m\to+\infty\). Specifically, we study the structural properties of any \(\rho^\ast\) and \(h^\ast\) for which there exists a sequence \(\left\{(\beta_n,\lambda_n,m_n)\right\}\) satisfying \(\beta_n,m_n\to+\infty\) and \(h^\ast =\lim_{n\to+\infty}h^\ast_{\beta_n,\lambda_n,m_n}\) and \(\rho^\ast =\lim_{n\to+\infty}\rho^\ast_{\beta_n,\lambda_n,m_n}\). The precise notion of these limits will be specified further below.

4 Main Results↩︎

In this section, we present our main results on the implicit bias of SGD when training wide two-layer ReLU networks for multivariate regression.

4.1 Implicit Bias Toward Piecewise Affine Functions↩︎

For a set of input data points \(\mathcal{X}=\left\{x_j\right\}_{j=1}^M \subset \mathbb{R}^d\), a linear dichotomy is an unordered partition \(\mathcal{X}=\mathcal{X}_+\sqcup \mathcal{X}_-\) such that \(\mathcal{X}_+\) and \(\mathcal{X}_-\) are strictly separable by an affine hyperplane; that is, there exist \(w\in \mathbb{R}^d\) and \(b\in \mathbb{R}\) satisfying \(x_j^\top w+b>0\) for \(x_j \in \mathcal{X}_+\) and \(x_j^\top w+b<0\) for \(x_j \in \mathcal{X}_-\). Note that a dichotomy is strict and unordered. It should be distinguished from a binary sign pattern on \(\mathcal{X}\), where the two classes carry specified labels and interchanging the labels gives a different sign pattern, and from a linear trichotomy of \(\mathcal{X}\), where the separation hyperplane is allowed to pass through some of the input points. Let \(\mathcal{P}(\mathcal{X})\) denote the total number of realizable linear dichotomies of \(\mathcal{X}\).

The following theorem shows that the function obtained by training a ReLU network with noisy SGD in the vanishing-temperature regime is piecewise affine.

Theorem 2, together with Theorem 1, shows that training an infinitely wide ReLU network with noisy SGD on training inputs \(\mathcal{X}\) results in a continuous piecewise affine function with at most \(2\mathcal{P}(\mathcal{X})-1\) kink hyperplanes, as the network width \(N\) tends to infinity and the temperature \(\beta^{-1}\) vanishes. Note that \(\mathcal{P}(\mathcal{X})\) depends only on the affine configuration of the input points \(\mathcal{X}\), and is independent of both the network width \(N\) and the ambient dimension \(d\). Hence, the bound on the linear-region complexity of the learned function depends only on the geometry of the training data.

By a classical result of [53], the bound ?? can be further estimated as follows: \[\label{eq:pwl-num-cover} 2\mathcal{P}(\mathcal{X})-1\leq 2 \sum_{i=0}^{d} \binom{M-1}{i}-1,\tag{8}\] where equality holds if and only if the training inputs \(\mathcal{X}\) are in general position, i.e., every subset of \(d\) or fewer points is linearly independent. In the non-general position case, \(\mathcal{P}(\mathcal{X})\) can be computed by via Zaslavsky’s Theorem [54] from the theory of hyperplane arrangements [55].

We emphasize that the regularization strength \(\bar \lambda = \lim_{n\to+\infty}\lambda_n\) can be an arbitrary positive real number. Thus, the strength of weight decay does not fundamentally affect the piecewise affine nature of the learned function. However, as shown in ?? , it controls the slope of the learned function: stronger regularization leads to flatter solutions and a larger tradeoff between data fitting and regularization. A numerical illustration of the role of the regularization strength is provided in Figure 5.

We reiterate that Theorem 2 applies to any function \(h^\ast\) that arises as the pointwise limit of a sequence \(h^\ast_{\beta_n,\lambda_n,m_n}\). Proposition [prop:accumu] shows that this class of accumulation points is non-empty and arises naturally: every sequence \((\beta_n,\lambda_n,m_n)\) satisfying Assumption 1, with \(m_n,\beta_n\to+\infty\) and \(\lambda_n\to \bar \lambda\), has a subsequence \((\beta_{n_k},\lambda_{n_k},m_{n_k})\) such that \(h^\ast_{\beta_{n_k},\lambda_{n_k},m_{n_k}}\) converges pointwise.

4.2 Feature Learning↩︎

Theorem 2 characterizes the implicit bias of noisy SGD in function space by showing that the algorithm is biased toward piecewise affine solution functions. We now examine the implicit bias in parameter space by studying the structure of the trained network parameters. This elucidates the feature-learning mechanism underlying the learning process.

Before stating our results, we emphasize that the piecewise affine structure of a solution function does not, in general, provide a tight characterization of the structure of the network parameters. For ReLU networks, the parameter fiber associated with a given function, i.e., the set of parameters realizing that function, may exhibit intricate geometric structure [56]–[61]. In particular, a piecewise affine solution function could arise from sparse output weights while leaving the input weights essentially unchanged, corresponding to feature selection rather than feature learning. Moreover, there exist continuous piecewise affine functions that cannot be represented exactly by any finite-width shallow ReLU network, such as \((u,v)\in \mathbb{R}^2 \mapsto \max(0,u,v)\). Therefore, to establish feature learning and effective width collapse, we investigate the stationary measure directly.

Consider the projection map \(\omega(a,w,b)=(w,b)\), and let \(\omega_\#\rho\) denote the push-forward on \(\mathbb{R}^{d+1}\) of a measure \(\rho \in \mathcal{P}(\mathbb{R}^{d+2})\). When \(\rho^\ast\) is the stationary measure of the mean-field PDE, the measure \(\omega_\#\rho^\ast\) is precisely the marginal measure of the input weights and biases of the trained network.

We now present our second main result, which establishes feature learning by showing that, after training, the input weights and biases align along finitely many data-dependent directions.

Theorem 3 shows that the input weights and biases of the trained network align along at most \(K \leq 2\mathcal{P}(\mathcal{X})-1\) directions. In particular, the learned function is representable by a shallow ReLU network with \(K\) hidden neurons: \[\label{eq:small-network} \int a(x^\top w+b)_+d\rho^\ast\; = \; \sum_{k=1}^K a_k (x^\top w_k +b_k)_+\;\text{with } a_k = \int_{\left\{(a,w,b)\colon (w,b)\in r_k\right\}} a \|(w,b)\| d\rho^\ast.\tag{9}\] Thus, the effective width of the network collapses to at most \(2\mathcal{P}(\mathcal{X})-1\).

We remark that ?? includes the degenerate case where the support is exactly the origin, i.e., \(\operatorname{supp}(q^\ast)=\left\{0\right\}\). In this case, the represented function is identically zero, and the only learned feature is the zero feature. This occurs, for example, when \(y_j=0\) for all \(j\), in which case the stationary measure is the Gaussian \(\rho^\ast(\theta)\propto \exp(-\frac{\beta \lambda}{2}\|\theta\|^2)\), which satisfies the Boltzmann fixed point condition 5 . Outside this degenerate case, whenever \(\operatorname{supp}(q^\ast) \neq \left\{0\right\}\), there is at least one aligned direction, and the network learns sparse, nonzero features. Empirically, we observe that after training, the input weights and biases often separate into two groups: one concentrated near the origin, whose contribution to the represented function is negligible, and another consisting of parameters with non-negligible norms that align along a few directions and constitute the learned nonzero features.

Theorem 3 applies to any measure \(\rho^\ast\) that arises as the \(2\)-Wasserstein limit of a sequence \(\rho^\ast_{\beta_n,\lambda_n,m_n}\). The assumption of \(2\)-Wasserstein convergence is used in the proof to obtain uniform integrability of the second moments (see Appendix 11). Establishing the existence of such a convergent sequence appears to be challenging in general. Nevertheless, Lemma 4 shows that any sequence \((\beta_n,\lambda_n,m_n)\) satisfying Assumption 1, with \(m_n,\beta_n\to+\infty\) and \(\lambda_n\to \bar\lambda>0\), admits a subsequence \((\beta_{n_k},\lambda_{n_k},m_{n_k})\) for which \(\rho^\ast_{\beta_{n_k},\lambda_{n_k},m_{n_k}}\) converges weakly.

4.3 Non-Redundancy of the Learned Features↩︎

We have shown that the input weights and biases of the trained network align along \(K\le 2\mathcal{P}(\mathcal{X})-1\) directions, each corresponding to an effective hidden neuron as in 9 . We now characterize these directions and, consequently, the structure of the learned features.

For the given set \(\mathcal{X}=\left\{x_j\right\}_{j=1}^M\), define a linear ternary activation pattern on \(\mathcal{X}\) as a map \(A\colon \mathcal{X}\to \left\{-1,0,1\right\}\) for which there exist \(w\in \mathbb{R}^d\) and \(b\in \mathbb{R}\) satisfying \(\mathrm{sgn}(x_j^\top w+b)=A(x_j)\) for \(x_j\in \mathcal{X}\). We distinguish two subclasses of ternary activation patterns: if \(A(x_j)\in\{\pm 1\}\) for every \(x_j\), we call \(A\) a binary activation pattern. Otherwise, if \(A(x_j)=0\) for at least one \(x_j\), we call \(A\) a strict ternary activation pattern. We further introduce the following partial order on activation patterns: \[\label{eq:partial-order} A' \preceq A \Longleftrightarrow A(x_j)\in \left\{0, A'(x_j)\right\},\,\forall j.\tag{10}\] Thus, \(A' \preceq A\) means that \(A\) is obtained from \(A'\) by replacing some (possibly none) of the \(\pm1\) labels of \(A'\) with zeros, while leaving all remaining labels unchanged. In this case, we say that \(A\) is a degeneration of \(A'\). Any ternary activation pattern is a degeneration of itself: \(A\preceq A\).

We present our third main result, which describes the structure of the learned features of the ReLU networks after SGD training.

Theorem 4 shows a strong non-redundancy of the learned features, illustrated in Figure 2. Item \((i)\) shows that each effective neuron of the trained network induces a distinct activation pattern on the training inputs. Moreover, the activation patterns induced by different effective neurons are incomparable with respect to the partial order \(\preceq\). Item \((ii)\) shows that every effective neuron activates at least one training point. Item \((iii)\) states that, from any binary activation pattern \(\hat{A}\) on \(\mathcal{X}\), one can obtain at most two learned ternary patterns \(A_k\) by turning some of the \(\pm1\) labels of \(\hat{A}\) into zeros. Theorem 4 can be directly translated into a description of the locations of the kink hyperplanes of the learned functions; the formal statement is provided in Corollary 1.

Identifying the exact aligned directions, or the exact location of the kink hyperplanes, is challenging for general training data. We obtain the following result for symmetric training data sets of size two.

Proposition [prop:toy] characterizes exactly the features extracted from simple symmetric data sets; see Figure 3 for an illustration. Interestingly, the solution function may or may not be globally linear depending on the norm of the training inputs. Another interesting phenomenon is that, the kink hyperplanes may or may not pass through the training points, which is in contrast to the linear spline interpolation solutions observed in different training regimes [16], [62] and in the optimal solutions [40], [44]. A similar observation was made by [9] in the univariate setting: for the specific training set \(\{(-10,2),(10,2)\}\), they showed that no function that is linear on the interval \([-10,10]\) can minimize the free energy. Consequently, the solution function must have at least one knot in the open interval \((-10,10)\). Here, we consider a more general setting and explicitly characterize the location of the kink hyperplanes.

We remark that characterizing the exact location of the kink hyperplanes for more general data requires a fine-grained understanding of the residuals of the learned predictor at the training points. We discuss this direction further in Section 6.5. A complete characterization remains an interesting and important problem for future work.

5 Discussion of the Main Results↩︎

5.0.0.1 Generalization

The Lipschitz bound ?? provides a standard route toward a generalization bound for the limiting solution function \(h^\ast\) in Theorem 2. Specifically, under common boundedness assumptions of the input and output domains, the class of functions satisfying ?? has controlled Rademacher complexity [63]. Standard Rademacher-complexity arguments then yield a generalization bound [64]. We remark that the bound ?? follows from the uniform second-moment bound on the stationary measures \(\rho^\ast_{\beta,\lambda,m}\) in Lemma 4, which is a multivariate adaptation of Lemma A.2 in the work of [9].

The bound on the effective width in Theorem 3 also allows one to control the network complexity, e.g., by bounding the VC-dimension via results in the work of [65]. However, such a bound would be insufficient for generalization. A first issue is that the bound \(2\mathcal{P}(\mathcal{X})-1\) grows with the cardinality of the training set \(M\). Moreover, controlling only the effective width (or the number of linear regions) is often too weak for obtaining non-vacuous generalization bounds, as piecewise linear functions with a bounded number of regions still constitute a large function class. A similar dependence on the sample size also appears in the bound of [9] for univariate networks; we note that no generalization result is provided there either.

[34] obtained a generalization guarantee for univariate shallow ReLU networks based solely on a bound on the number of linear regions of the learned function. This is possible for two reasons: first, they considered a classification problem with the \(0\)-\(1\) loss, where the prediction error depends only on the sign of the network output, not on its magnitude. In this setting, bounding the number of linear regions, and hence the number of sign changes of the prediction function, is sufficient to control the complexity of the function class relevant to the \(0\)-\(1\) loss; second, they assumed that the labels are generated by a teacher network of width \(r\), and by analyzing a KKT condition satisfied by the trained student network, they derived a bound on the number of linear regions that depends only on \(r\), not on the sample size. The student-teacher setting is important, since fixed-size ReLU networks can have poor empirical performance on random labels even under the \(0\)-\(1\) loss [66].

5.0.0.2 The Role of Weight Decay

We considered a sequence of regularization parameters \(\lambda_n>0\) converging to an arbitrary limit \(\bar \lambda>0\). The positivity of each \(\lambda_n\) is needed to ensure the applicability of Theorem 1, in particular the validity of the Boltzmann fixed point condition 5 . As observed by [5], when \(\lambda=0\) and the neuron map \(\sigma\) is bounded, the right-hand side of 5 fails to be normalizable. The positivity assumption on \(\bar \lambda\) is needed in the analysis of the limiting support of the measure \(\rho^\ast_{\beta,\lambda,m}\) (see Proposition [prop:shrink]).

We note that the particular value of \(\bar \lambda\) has no bearing on the piecewise-affine nature of the solution function, the alignment of the learned input weights and biases, or the non-redundancy property of the learned features. Nevertheless, it controls the Lipschitz constant of the solution function; see ?? and Figure 5.

5.0.0.3 The Role of Stochasticity in SGD

The stochasticity in SGD, including both data sampling and injected Gaussian noise, is needed in the mean-field analysis in Theorem 1. Several alternative training procedures have also been studied in the mean-field regime, such as mini-batch and full-batch training, training without noise, training with Gaussian noise under different scalings, and training with heavy-tailed noise [5], [6], [8], [67]–[69]. In these settings, however, either convergence of the mean-field PDE dynamics in time to a stationary measure has not been established, or the stationary measures lack a sufficiently explicit characterization. As a result, analyzing the corresponding implicit bias remains challenging. Nevertheless, neither the specific data sampling scheme nor the precise noise model enters our analysis. The proofs rely only on the Boltzmann fixed point condition 5 and on structural properties of the ReLU neuron.

5.0.0.4 The Role of Initialization and Scaling

Notice that Theorem 1 imposes only mild assumptions on the initialization distribution \(\rho_0\): namely, that \(\rho_0\in \mathcal{P}_2^a\) and is sub-Gaussian. Consequently, our results apply to any initialization satisfying these conditions. This broad applicability stems from the injected Gaussian noise in SGD, or equivalently from the entropy regularization appearing in the free energy. The entropy term makes the free energy strictly convex, thereby ensuring the existence of a unique stationary measure.

We consider wide neural networks in the mean-field regime, using the parametrization with one-over-width scaling in 1 and an initialization distribution that is independent of the network width. Alternative choices of the parametrization and initialization scaling give rise to different training regimes, with different training dynamics and implicit biases under gradient-based optimization [10], [11], [16].

5.0.0.5 Comparison with [9]

[9] studied univariate network with \(m\)-truncated, \(\tau\)-smoothed ReLU neurons, trained by noisy SGD with temperature \(\beta^{-1}\) and weight decay \(\lambda\) on \(M\) data points. They showed that, as \(\beta,\tau,m\to+\infty\), the learned function is approximated by piecewise linear functions with \(O(M)\) knots. Theorem 2 extends this result to the multivariate case \(d\geq1\), where the bound on the number of kink hyperplanes has order \(O(\sum_{i=0}^{d} \binom{M-1}{i})\). Regarding the configuration of the knots, [9] showed that there are at most three knots between any two consecutive training points and, when three knots occur, two must coincide with the endpoints of the interval, while the third lies strictly in its interior. This appears as a special case of Corollary 1, which is a direct consequence of Theorem 4 (see Appendix 16 for details). Moreover, Corollary 1 yields an additional consequence in the univariate setting: there can be at most two knots in \((-\infty, x_1]\cup [x_M,+\infty)\), where \(x_1=\min_j x_j\) and \(x_M=\max_j x_j\). If two knots are present, they must occur precisely at the boundary data points \(x_1\) and \(x_M\). To our knowledge, this observation is not contained in the analysis of [9].

Beyond extending the above results to general input dimensions, a key advance over [9] is that we investigate the implicit bias in parameter space. Specifically, Theorem 3 shows that the learned input weights and biases align along finitely many directions, leading to an effective collapse of the network width. Theorem 4 further characterizes the structure of these directions and the resulting learned features. As discussed in Section 4.2, these results are nontrivial and do not follow from the piecewise-affine structure of the represented function alone.

Compared with [9], the present work has the following limitations. The piecewise-linearity result of [9] applies to \(\lim_n \lambda_n\ge 0\) including the regime of vanishing regularization, whereas our analysis requires \(\bar \lambda=\lim_n \lambda_n>0\). As discussed earlier, the positivity of \(\bar \lambda\) is needed to analyze the limiting support of the stationary measure \(\rho^\ast_{\beta,\lambda,m}\). The support analysis is a key ingredient in our study of the parameter-space implicit bias and, to our knowledge, has no counterpart in the work of [9], where the support of the limiting measure is not investigated.

Additionally, to establish the piecewise-affine structure of the solution function, both [9] and our work rely on a bound on the second derivative of the solution function for \(x\) outside of a failure set \(U\). This bound vanishes as \(\beta,m\to+\infty\) and the set \(U\) shrinks to a union of points in the univariate case or hyperplanes in the multivariate case. In the univariate setting, [9] further derives a non-asymptotic bound on the Lebesgue measure of \(U\). We do not obtain an analogous bound for \(d>1\). In our analysis, the failure set \(U\) consists of \(x\) that violates the Quadratic Positivity Condition. This condition requires that, in the space of hyperplanes, the set \(\hat{G}_x\) of all hyperplanes passing through \(x\) be disjoint from another set \(\hat{\Omega}^{\beta,\lambda,m}\) (see Sections 6.3–6.4 for details). Although we obtained a non-asymptotic bound on the size of \(\hat{\Omega}^{\beta,\lambda,m}\) (see Appendix 10.3), converting this information into a quantitative estimate for the measure of the failure set in the input domain appears challenging.

5.0.0.6 Comparison with [38], [39]

[38], [39] studied the minimal total variation norm (TV norm) interpolation problem: \[\label{eq:del} \inf_{\mu\in M(\mathbb{S}^{d})} \|\mu \|_{\mathrm{TV}} \quad \mathrm{s.t. }\;\int (w^\top x_j+b)_+ d\mu(w,b) = y_j,\quad j=1,\ldots,M,\tag{11}\] where \(M(\mathbb{S}^{d})\) is the space of signed Radon measures on \(\mathbb{S}^{d}\). They proved that all minimizers of 11 are atomic measures, of the form \(\rho^\ast = \sum_{k=1}^{K'}c_i \delta_{w_i}\), and thus correspond to ReLU networks with at most \(K'\) hidden neurons. As discussed in Section 2, the key difference between problem 11 and the free-energy-minimization problem considered in this work is their connection to training dynamics: the latter precisely characterizes the implicit bias of noisy SGD, whereas there is no general guarantee that the minimizer of problem 11 is reached by the training procedure. Moreover, due to the entropy regularizer, the minimizer of the free energy at any finite temperature \(\beta^{-1}>0\) must be absolutely continuous. By contrast, atomic measures belong to the feasible set of 11 and are naturally promoted by the TV regularizer.

[38] showed that \(K'\leq 2\mathcal{P}(\mathcal{X})\) (see the proof of Theorem 4.2 therein), which exceeds our bound ?? by one. [39] showed that \(K' \leq \max_{k=0,\ldots,M} \binom{M}{k}\Xi(M-k,d+1-k)\), where \(\Xi(n,d)\) is the number of regions in \(\mathbb{R}^d\) produced by an arrangement of \(n\) central hyperplanes. In Appendix 14 we show that, for generic data, this bound is strictly larger than ours, and that in the regime \(2d+1 \leq M\), our bound improves upon it by a factor of at least \(2^d/(d+1)^2\).

6 Proof of the Main Results↩︎

In this section, we outline the proofs of our main results. In our analysis, we mostly focus on the input weights and biases of the ReLU network, as they define the hidden features and the kink hyperplanes of the output function. Thus, throughout this section, by parameter space we mean the space of input weights and biases \(\mathbb{R}^{d+1}\).

In Section 6.1, we introduce the notation and a decomposition of the parameter space induced by the training data. In Section 6.2, we derive a bound on the limit of the Hessian of the solution function \(h^\ast_{\beta,\lambda,m,\tau}(x)\) as \(\tau\to+\infty\). In Section 6.3, we introduce the Quadratic Positivity Condition at a test location \(x\). This condition requires, in the parameter space, that the set of all hyperplanes passing through \(x\) be disjoint from a cluster set \(\Omega^{\beta,\lambda,m}\), defined by the non-positive sets of certain quadratic functions appearing in the Hessian bound. For every \(x\) satisfying this condition, we obtain a further bound on the limiting Hessian that vanishes as \(\beta,m\to+\infty\). In Section 6.4, we show that the cluster set \(\Omega^{\beta,\lambda,m}\) shrinks to a limiting set \(\Omega^\ast\), given by a union of \(K\le 2\mathcal{P}(\mathcal{X})-1\) rays in \(\mathbb{R}^{d+1}\). These rays define \(K\) hyperplanes in input space. Based on this, we then show that (i) for every \(x\) at positive distance from these \(K\) hyperplanes, the Hessian at \(x\) vanishes as \(\beta,m\to+\infty\); and (ii) the marginal measure of input weights and biases has vanishing mass outside these \(K\) rays. Finally, in Section 6.5, we investigate the geometry of the limiting set \(\Omega^\ast\).

Our proof strategy builds on the analysis of [9] for univariate networks. Their argument has two main components. First, they derive a bound on the Hessian at a test location \(x\) that vanishes as \(\beta,m\to+\infty\), provided that an associated quadratic function \(f_{\beta,m}\) is positive at \(x\). Second, they show that the non-positive set of \(f_{\beta,m}\) has vanishing Lebesgue measure as \(\beta,m\to+\infty\), and hence shrinks to a finite set of points. In their analysis, the quadratic function \(f_{\beta,m}\) is analyzed on the input domain and plays a central role. As noted by the authors, extending these arguments beyond one dimension presents two main obstacles. First, there is no simple analogue of the interval decomposition of the input domain: in \(\mathbb{R}^1\), the training inputs partition the input domain into intervals, on each of which the activation pattern of every ReLU neuron is fixed. Second, whereas in one dimension the non-positive set of a quadratic function is the union of at most two intervals, in higher dimensions the corresponding quadric sublevel sets can exhibit considerably more intricate geometry, making their concentration difficult to analyze. We overcome both obstacles by shifting the analysis from the input space to the parameter space, where the training data induce a natural decomposition. We identify the relevant quadratic functions on the parameter space and analyze the shrinkage of their non-positive sets.

In the sequel, we let \(\alpha=(\beta,\lambda,m)\) and, for two quantities \(A,B\), we write \(A\lesssim B\) if there exists a constant \(C\) independent of \((\beta,\lambda,m,\tau,x)\), such that \(A \leq C B\).

6.1 Subdivision of Parameter Space↩︎

We begin by introducing a natural decomposition of the parameter space induced by the training data, together with the notation needed for our analysis.

For a single ReLU neuron \(a(w^\top x + b)_+\), the non-linearity occurs along the hyperplane \(\{x \in \mathbb{R}^d \colon w^\top x + b = 0\}\). This defines a natural mapping from the parameter space to the set of affine hyperplanes in the input space. Conversely, for any \(x\in \mathbb{R}^d\), let \(G_{x}= \left\{(w,b)\in \mathbb{R}^{d+1}\colon w^\top x + b=0\right\}\), which is the set of parameters \((w,b)\) whose associated input-space hyperplane passes through \(x\). Given a training set \(\mathcal{X}=\{x_j\}_{j=1}^M\), consider the arrangement of hyperplanes \(\{G_{x_j}\}_{j=1}^M\subset \mathbb{R}^{d+1}\). Let \(\left\{U_k\right\}\) denote the sectors of the arrangement, that is, the open connected components, of \(\mathbb{R}^{d+1}\setminus \cup_j G_{x_j}\). Each sector \(U_k\) corresponds to a unique binary sign pattern on \(\mathcal{X}\) and determines an index set \[J_k=\left\{j \in \left\{1,\ldots,M\right\} \colon \mathrm{sgn}(w^\top x_j + b)=1,\, \forall (w,b)\in U_k\right\}.\] Thus, the sectors are in one-to-one correspondence with the binary sign patterns on \(\mathcal{X}\) realizable by affine functions. The number of such binary sign patterns is exactly twice the number of linear dichotomies. Hence, there are precisely \(2\mathcal{P}(\mathcal{X})\) sectors in total. For analytical convenience, we work with their closures and write \(O_k = \mathrm{cl}(U_k)\) for \(k=1,\ldots,2\mathcal{P}(\mathcal{X})\).

Fix a point \(x \in \mathbb{R}^d\) at which the Hessian of the solution function is to be evaluated. The hyperplanes \(\{G_{x_j}\}_{j=1}^M\) induce a decomposition of \(G_x\), \(G_x = \bigcup_k (G_x \cap O_k )\). Each open sector \(G_x \cap U_k\) consists of those parameters that define a hyperplane passing through the evaluation point \(x\) while inducing the binary sign pattern with index set \(J_k\) on the training set. The set \(G_x\) can be naturally identified with \(\mathbb{R}^d\) via the map \(w\mapsto (w,-w^\top x)\). Then, for \(k=1,\ldots,2\mathcal{P}(\mathcal{X})\), the closed sector \(G_x\cap O_k\) can be identified with \[O_k^w=\left\{w\in\mathbb{R}^d \colon (w,-w^\top x)\in O_k\right\}.\]

6.2 Zero-Smoothing Limit↩︎

Next, we control the Hessian of the solution function in the zero-smoothing limit \(\tau\to+\infty\). By the construction of \(\sigma^{m,\tau}\) in Appendix 9, the limits \((u)_+^m=\lim_{\tau \to+\infty} (u)_+^{m,\tau}\) and \(v^m=\lim_{\tau \to+\infty} v^{m,\tau}\) exist for every \(m>0\); in particular, \[\label{eq:vm} v^m=\lim_{\tau \to+\infty} v^{m,\tau}= \begin{cases} v, & \|v\|\leq m\\ \frac{m}{\|v\|} v, &\|v\|>m. \end{cases}\tag{12}\] Let \(\sigma^m\) denote the \(m\)-truncated ReLU neuron in the zero-smoothing limit: \[\sigma^m(x,\theta)=\lim_{\tau\to+\infty}\sigma^{m,\tau}(x,\theta)=a^m(x^\top w^m +b)_+^m .\] Since the maps \(u\mapsto (u)_+^{m,\tau}\) and \(v\mapsto v^{m,\tau}\) are bounded and have bounded first derivatives with the bounds independent of \(\tau\), \(\sigma^m\) is bounded and Lipschitz continuous for every \(m>0\).

By results of [5], the free energy \(F_{\sigma^{m}}^{\beta,\lambda}\) associated with the \(m\)-truncated ReLU neuron \(\sigma^m\), with regularization \(\lambda>0\) and temperature \(\beta^{-1}>0\), admits a unique minimizer, which we denote by \(\rho^\ast_\alpha\). Define \[h^\ast_{\alpha}(\cdot)=h_{\sigma^m}(\cdot,\rho^\ast_{\alpha}), \qquad R_j^\alpha = \frac{1}{M} \Big( h^\ast_{\alpha}(x_j)- y_j \Big),\, j=1,\ldots,M.\] Thus, \(h_\alpha^\ast\) is the solution function associated with the \(m\)-truncated neuron \(\sigma^{m}\) and \(R_j^\alpha\) is the residual of \(h^\ast_{\alpha}\) at the \(j\)th data point, divided by \(M\).

A priori, the limit \(\lim_{\tau\to+\infty} h^\ast_{\beta,\lambda,m,\tau}\) need not coincide with \(h^\ast_{\alpha}\), since \(h^\ast_{\beta,\lambda,m,\tau}\) is obtained by first minimizing the free energy associated with the smoothed and truncated neuron \(\sigma^{m,\tau}\) and then taking \(\tau\to+\infty\), whereas \(h^\ast_\alpha\) is obtained by minimizing the free energy associated directly with the limiting neuron \(\sigma^m\). The following result shows that, as \(\tau\to+\infty\), not only \(h^\ast_{\beta,\lambda,m,\tau}\) but also its Hessian converge to \(h^\ast_\alpha\) and its associated Hessian, respectively. This proves Lemma [prop:tau]. We also bound the norm of the Hessian in the zero-smoothing limit. The proof is provided in Appendix 10.1.

Thus, the Hessian is upper bounded by three terms. The first two vanish as \(m,\beta\to+\infty\). The third is more complicated and depends on \(\mathcal{I}^\alpha(x)\).

6.3 Quadratic Positivity Condition↩︎

We introduce a condition for controlling the term \(\mathcal{I}^\alpha(x)\) in the right-hand side of 13 and thus controlling the Hessian. To motivate the condition, recall from 12 that \(w^m=\gamma_w w\) for some \(\gamma_w\in [0,1]\). Observe that if \[\label{eq:hessian-condition-2} \|w\|^2+|x^\top w|^2-|(Q_k^\alpha)^\top w|^2 \geq \delta \|w\|^2, \quad \forall w \in O_k^w, \;\forall k,\tag{14}\] for some \(\delta \in (0,1]\), then, \[\label{eq:okw} \begin{align} \|w\|^2+|x^\top w^m|^2 -|(Q_k^\alpha)^\top w^m|^2 &=(1-\gamma_w^2)\|w\|^2 + \gamma_w^2(\|w\|^2+|x^\top w|^2-|(Q_k^\alpha)^\top w|^2)\\ & \geq (1-\gamma^2_w(1-\delta))\|w\|^2\\ & \geq \delta\|w\|^2. \end{align}\tag{15}\] Under this condition, each integral in \(\mathcal{I}^\alpha(x)\) can be controlled by the third moment of a Gaussian distribution as follows: \[\begin{align} &\int_{O_k^w} \|w\|^2 |(Q^\alpha_k)^\top w^m| \exp\left\{-\frac{\lambda\beta}{2}\Big( \|w\|^2+|x^\top w^m|^2-|(Q^\alpha_k)^\top w^m|^2\Big)\right\} dw \\ \leq & \sqrt{1+\|x\|^2}\int_{\mathbb{R}^d} \|w\|^3 \exp\left\{-\frac{\lambda\beta \delta}{2}\|w\|^2\right\} dw\\ \lesssim& \sqrt{1+\|x\|^2} (\lambda\beta\delta)^{-\frac{d+3}{2}}. \end{align}\] Then the Hessian can be bounded as \[\|\nabla_x^2 h^\ast_{\alpha}(x)\| \lesssim m^{-1}(1+\lambda^{-1}) +\beta^{-1}\lambda^{-\frac{d+6}{4}}(1+\sqrt{\lambda})^{\frac{d+2}{2}} \Big( 1 + \delta^{-\frac{d+3}{2}} \sqrt{1+\|x\|^2}\Big),\] where the right-hand side vanishes as \(m,\beta\to+\infty\), provided that \(\delta\) is lower bounded.

We reformulate condition 14 in a more convenient form. For \(k=1,\ldots,2\mathcal{P}(\mathcal{X})\), define a quadratic function \[\label{eq:f} f^\alpha_k(w,b) = \|w\|^2+b^2 - \Big( \frac{1}{\lambda}\sum_{j\in J_k} R^\alpha_j (x_j^\top w +b) \Big)^2,\tag{16}\] with the convention that \(f_k^\alpha(w,b)=\|w\|^2+b^2\) if \(J_k=\varnothing\). Note that \(f_k^\alpha\) is 2-homogeneous and that the left-hand side of 14 is precisely \(f^\alpha_k(w,-w^\top x)\). Thus, 14 is satisfied if \[\label{eq:qpc} \min_{k\in \left\{1,\ldots,2\mathcal{P}(\mathcal{X})\right\}} \inf_{\substack{(w,b)\in G_x\cap O_k\\ \|w\|^2+b^2=1}} f_k^\alpha(w,b)>0,\tag{17}\] with the convention that \(\inf_{\substack{(w,b)\in S}} f(w,b)=+\infty\) if \(S=\varnothing\). We refer to 17 as the Quadratic Positivity Condition, as it requires the positivity of \(f_k^\alpha\) on the intersection of the hyperplane \(G_x\) and the closed sector \(O_k\) for every \(k\).

We have the following result bounding the integral in the right-hand side of 13 and thus the Hessian of the solution function at \(x\), provided that \(x\) satisfies the Quadratic Positivity Condition. The proof is provided in Appendix 10.2.

6.4 Shrinkage of Cluster Set↩︎

According to the bound 19 , the Hessian vanishes as \(m,\beta\to+\infty\) and \(\lambda\to\bar \lambda>0\), provided that the margin \(\delta\) remains uniformly bounded away from zero throughout this limiting process. We now examine when \(\delta\) admits such a lower bound.

To this end, we derive a geometric description of the Quadratic Positivity Condition. Let \[P_k^\alpha = \frac{1}{\lambda}\sum_{j\in J_k} R^\alpha_j \begin{pmatrix} x_j\\1 \end{pmatrix} \in \mathbb{R}^{d+1},\] with the convention that \(P_k^\alpha=0\) if \(J_k=\varnothing\). Writing \(z=(w,b)\), we can express \(f_k^\alpha\) as the quadratic form \[f_k^\alpha(z)= z^\top \left(I_{d+1} -P_k^\alpha (P_k^\alpha)^\top \right) z = z^\top \Lambda_k^\alpha z.\] The coefficient matrix \(\Lambda_k^\alpha\) has eigenvalues \(1\) with multiplicity \(d\) and \(1-\|P_k^\alpha\|^2\). Thus, if we define the non-positive set of \(f_k^\alpha\) by \[\Omega_k^\alpha = \left\{(w,b)\in\mathbb{R}^{d+1} \colon f_k^\alpha (w,b)\leq 0\right\},\] then \(\Omega_k^\alpha\) is a quadratic double cone when \(\|P_k^\alpha\|^2>1\), a line when \(\|P_k^\alpha\|^2=1\), and the singleton \(\left\{(0,0)\right\}\) when \(\|P_k^\alpha\|^2<1\). Here and throughout, we call a set \(S\subset \mathbb{R}^{d+1}\) a cone if \(v\in S\) implies \(c v\in S\) for all \(c \geq 0\). Notably, 17 is equivalent to requiring that the cones \(G_x\cap O_k\) and \(\Omega_k^\alpha\) intersect only at the origin. This condition can be expressed on the unit sphere. For any cone \(S\subset \mathbb{R}^{d+1}\), let \(\hat{S} = S\cap \mathbb{S}^d\) denote its spherical cross-section. Since a cone is fully determined by its spherical cross-section, 17 is equivalent to \[(\hat{G}_x \cap \hat{O}_k ) \cap \hat{\Omega}_k^\alpha = \varnothing, \quad \forall k,\] illustrated in Figure 4, or, equivalently, \[\hat{G}_x \cap \hat{\Omega}^\alpha = \varnothing, \quad \text{where} \quad \hat{\Omega}^\alpha = \bigcup_{k=1}^{2\mathcal{P}(\mathcal{X})} \hat{O}_k \cap \hat{\Omega}^\alpha_k.\]

We refer to the set \(\hat{\Omega}^\alpha\) as the cluster set, in analogy to the terminology of [9]. Next we show that, under a suitable sequence \(\alpha_n=(\beta_n,\lambda_n,m_n)\), the cluster set \(\hat{\Omega}^{\alpha_n}\) shrinks to at most \(2\mathcal{P}(\mathcal{X})-1\) fixed points on the sphere \(\mathbb{S}^d\), as \(n\to+\infty\). Consequently, for any \(x\in\mathbb{R}^d\) such that \(\hat{G}_x\) avoids these limiting points, we can obtain a uniform positive lower bound on the margin \(\delta\). The following result formalizes this shrinkage of the cluster set. Its proof is provided in Appendix 10.3.

Notice that the limit \(\lim_{n\to +\infty}P_k^\alpha\) exists for every \(k\), as both limits, \(\lim_{n\to+\infty} \lambda_n\) and \(\lim_{n\to+\infty}R^{\alpha_n}_j\), exist by our assumptions.

Based on Proposition [prop:qpc] and Proposition [prop:shrink], we obtain the following result, thereby proving Theorem 2. The proof is given in Appendix 10.4.

The shrinkage of the cluster set also allows us to analyze the support of the stationary measure \(\rho^\ast_{\alpha}\). Specifically, for any set \(B\subset \mathbb{R}^{d+1}\) such that the projection of \(\mathrm{cl}(B)\) onto the sphere is disjoint from \(\hat{\Omega}^\ast\), we show that \[\lim_{n\to+\infty}\mathbb{E}_{\theta\sim\rho^\ast_{\alpha_n}}\left[\|(w,b)\|^2 \cdot \mathbb{I}_{(w,b)\in B}\right]=0.\] That is, the \(\|(w,b)\|^2\)-weighted mass of \(\rho^\ast_{\alpha_n}\) over \(\mathbb{R}\times B\) vanishes as \(n\to+\infty\). This directly gives the desired results in Theorem 3. The formal proof is given in Appendix 11.

6.5 Position of Limiting Cluster Set Relative to Sphere Stratification↩︎

In view of the previous subsection, the limiting set \(\hat{\Omega}^\ast\) determines the aligned directions in parameter space and the kink hyperplanes in input space. We now study the location of the points in \(\hat{\Omega}^\ast=\left\{(w_k,b_k)\right\}\) relative to a stratification of the sphere \(\mathbb{S}^d\) induced by ternary activation patterns on the training inputs.

Recall from Section 4.3 the definition of linear ternary activation patterns on the training set. Each ternary activation pattern \(A\colon \mathcal{X}\to\left\{-1,0,1\right\}\) defines a stratum on the sphere \(\mathbb{S}^d\): \[V_A = \left\{(w,b)\in \mathbb{S}^d \colon \mathrm{sgn}(w^\top x_j +b)=A(x_j),\;\forall j \right\}.\] Let \(\mathcal{V}=\{V_A\}_A\), where \(A\) ranges over all linear ternary activation patterns realizable on \(\mathcal{X}\). The following basic result shows that \(\mathcal{V}\) forms a stratification of the sphere \(\mathbb{S}^d\) and the strata have simple geometry. Its proof is given in Appendix 12.

Notably, all the notions related to ternary activation patterns introduced in Section 4.3 can be translated into geometric language. To know the activation pattern induced by the hyperplane \(\pi_k=\{x\colon w_k^\top x+b_k=0\}\), it suffices to know which stratum of \(\mathcal{V}\) contains the point \((w_k,b_k)\). A ternary activation pattern \(A\) is a binary sign pattern if and only if \(V_A\) is a \(d\)-dimensional stratum, i.e., has the same dimension as the sphere \(\mathbb{S}^d\), and is a strict ternary sign pattern if and only if \(V_{A}\) has positive codimension on \(\mathbb{S}^d\). For the partial order on ternary sign patterns defined in 10 , it is straightforward to verify that \[A' \preceq A \Longleftrightarrow \mathrm{cl}(V_{A'}) \supset V_A.\] In addition, for a given binary sign pattern \(\hat{A}\), the set \(\mathcal{D}(\hat{A})\) defined in ?? can be rewritten as \[\mathcal{D}(\hat{A}) = \left\{A_k \colon (w_k,b_k)\in \mathrm{cl}(V_{\hat{A}}),k=1,\ldots, K\right\}.\]

Exploiting the geometry of the limiting set \(\hat{\Omega}^\ast\), together with the structure of the strata in Lemma 2, we obtain the following result, which provides a geometric proof of Theorem 4. Its proof is provided in Appendix 12. We write \(\partial_{\mathrm{cl}(S)}S\) for the relative boundary of \(S\) in its closure for any \(S\subset \mathbb{S}^d\).

Finally, we briefly outline the proof of Proposition [prop:toy]. As indicated by Theorem 5, the vectors \[P^\ast_k = \frac{1}{\bar \lambda M} \sum_{j\in J_k} \left(h^\ast(x_j)-y_j\right)\begin{pmatrix} x_j\\1 \end{pmatrix},\quad k=1,\ldots,2\mathcal{P}(\mathcal{X}),\] characterize the set \(\hat{\Omega}^\ast\) and hence the locations of the kink hyperplanes \(\{\pi_k\}_{k=1}^K\). For the specific data set \(\mathcal{X}=\{(x_1,y_1),(x_2=-x_1,y_2=y_1)\}\), we have (i) \(h^\ast\) is an even function and \(h^\ast(x_1)-y_1=h^\ast(x_2)-y_2\); and (ii) there are only four possible binary sign patterns: \(J_1=\left\{1,2\right\},J_2=\left\{1\right\},J_3=\left\{2\right\},J_4=\varnothing\). These observations allow us to explicitly compute \(P^\ast_k\) for \(k=1,\ldots,4\), thereby identifying the locations of the hyperplanes. The case of odd labels follows a similar strategy. The formal proof is provided in Appendix 13. For general training data, determining \(P^\ast_k\) remains challenging, in particular if one relies solely on the Boltzmann fixed point condition. We expect that additional information about the free-energy minimizer is needed.

7 Experiments↩︎

In Figure 1 and Figure 2, we consider a two-layer ReLU network of width \(N=1000\) and initialize the neuron parameters \((a,w,b)\) as i.i.d.samples from the Gaussian \(\mathcal{N}(0, \nu^2 I_{d+2})\) with \(\nu=10^{-3}\). On a training data set of size \(M=9\) with two-dimensional inputs, we run the SGD algorithm 2 with step size \(s_{N,k} \equiv 10^{-3}\), weight decay regularization strength \(\lambda=10^{-2}\) and no Gaussian noise, \(\beta^{-1}=0\), for \(2\times 10^6\) iterations. The learned features are identified as follows: (i) among all trained input weights and biases \({(w^i,b^i)}_{i=1}^N\), keep only those with norm larger than \(10^{-8}\) and normalize them onto the unit sphere; and (ii) apply DBSCAN [70], a standard clustering method, to the normalized points under the spherical distance: \(d_\angle(z,z')=\arccos(z^\top z')\) for \(z,z'\in \mathbb{S}^d\). In our experiment, the resulting clusters had maximal within-cluster spherical distance \(7.9\times 10^{-4}\), approximately \(0.045^\circ\). These results are consistent with Theorem 3 and Theorem 4.

In Figure 3, we report the result of training a two-layer ReLU network on a univariate data set of size \(M=2\) via SGD for \(10^6\) iterations. The network width, initialization, step size, regularization, and noise level are the same as in the previous experiment. The result is consistent with Proposition [prop:toy].

Figure 5 and Figure 6 illustrate the roles of weight decay and Gaussian noise in SGD, respectively. In these figures, a two-layer ReLU network is trained on a two-dimensional data set via SGD for \(10^6\) iterations. The network width, initialization, and baseline SGD parameters are the same as before. The results are consistent with Theorem 2.

8 Conclusion↩︎

We presented an analysis of SGD training of wide shallow ReLU networks. We showed in the mean-field limit that the learned function is piecewise linear. At convergence, the input weights and biases of the network align along a finite number of directions depending on the training data, so the network learns a finite collection of features and its effective width collapses. These directions satisfy a non-redundancy property: distinct directions induce distinct ternary linear sign patterns on the training inputs.

These results contribute to the theoretical understanding of the implicit bias and feature-learning mechanisms of gradient-based training in nonlinear neural networks. While it is widely believed that gradient-based optimization can discover useful feature representations, a rigorous mathematical understanding of this phenomenon has remained limited. Previous theoretical results establishing the emergence of sparse or structured features have either relied on specialized assumptions on the training data or initialization, or have characterized optimal solutions of associated variational problems. In contrast, our analysis shows that an effectively finite width representation emerges directly from the training dynamics in a broad mean-field setting, and relates its structure to the combinatorial geometry of the training data.

8.0.0.1 Limitations

Our analysis is formulated within the mean-field framework of [5] and therefore considers two-layer networks in the infinite-width limit trained by noisy SGD with weight decay. Rather than analyzing ReLU neurons directly, we study bounded, smooth approximations and subsequently pass to the ReLU limit. This approximation is motivated by the current state of the theory. To our knowledge, a direct initialization-to-convergence analysis for mean-field training of ReLU networks has not yet been established.

8.0.0.2 Future work

Several interesting directions remain open. A natural next step is to obtain finer-grained descriptions of the learned features, including characterizations of the precise locations of the kink hyperplanes beyond the special classes of training data considered in Proposition [prop:toy]. Another important question is whether the bound on the number of learned features is tight. Our numerical experiments, together with those of [9], suggest that this bound may admit further improvement, possibly depending on the specific structure of the data. It would also be of interest to extend the analysis to other optimization procedures (e.g., gradient descent, mini-batch SGD, and preconditioned SGD), classification losses, different activation functions, and more complex architectures such as deep networks or attention-based models.

Reproducibility Statement↩︎

Code to reproduce our experiments is available at https://github.com/shuangliang15/effective-width-collapse.

Appendix↩︎

The appendix is organized into the following sections.

• Appendix 9: Approximation of the ReLU Neuron

• Appendix 10: Proof of Theorem 2

• Appendix 11: Proof of Theorem 3

• Appendix 12: Proof of Theorem 4

• Appendix 13: Proof of Proposition [prop:toy]

• Appendix 14: Comparison of the Bounds on Effective Width

• Appendix 15: Existence of Accumulation Points

• Appendix 16: Non-redundancy of Kink Hyperplane Arrangement

9 Approximation of the ReLU Neuron↩︎

In this section, we provide a detailed construction of the \(m\)-truncated, \(\tau\)-smoothed ReLU neuron \(\sigma^{m,\tau}\).

9.1 Truncated and Smoothed ReLU Neuron↩︎

For approximating the ReLU activation function, let \[(u)_+^{m,\tau} = \begin{cases} S_\tau(u)=\log(1+e^{\tau u})/\tau,& u < S_\tau^{-1}(m^2)\\ \phi_{m,\tau}(u), & u\ge S_\tau^{-1}(m^2), \end{cases}\quad u\in \mathbb{R}.\] The tail is chosen so that \((\cdot)_+^{m,\tau}\in C^4(\mathbb{R})\) and satisfies \[\|\phi_{m,\tau}\|_\infty\le 2m^2,\qquad \|\phi_{m,\tau}'\|_\infty \le 1,\qquad \|\phi_{m,\tau}''\|_\infty \le m^{-2},\] with bounded third and fourth derivatives. Moreover, as \(\tau\to\infty\), \((\cdot)_+^{m,\tau}\) converges to a truncated ReLU \((\cdot)_+^m\), defined as follows: \[(u)_+^m=\begin{cases} (u)_+,& u< m^2\\ \phi_m(u),& u\geq m^2, \end{cases}\quad u\in \mathbb{R},\] which satisfies that \((\cdot )_+^m\in C^4((0,+\infty))\) and \(\phi_{m,\tau}''\) converges to \(\phi_m''\) uniformly on any compact subset of \((m^2,+\infty)\). An explicit construction of \(\{(\cdot)_+^{m,\tau}\}\) is given in the next subsection.

For approximating the identity map, let \[v^{m,\tau} = \begin{cases} v,& \|v\|< m-\tau^{-1}\\ \psi_{m,\tau}(v), &\|v\|\ge m-\tau^{-1}, \end{cases}\quad v\in \mathbb{R}^n.\] The tail function \(\psi_{m,\tau}\) is selected such that \((\cdot)^{m,\tau}\in C^4(\mathbb{R}^n,\mathbb{R}^n)\), is an odd function, and satisfies that \(\|\psi_{m,\tau}(u)\|\leq m, \|\nabla \psi_{m,\tau}(u)\|\leq 1\), and \(\|\nabla^{(k)} \psi_{m,\tau}\|\) are bounded for \(k=2,3,4\). Additionally, we require that as \(\tau\to+\infty\), \((\cdot)^{m,\tau}\) converges to the truncated identity map \((\cdot)^m\), defined as follows \[v^m=\begin{cases} v,&\|v\|<m\\ m\cdot v/\|v\|, &\|v\|\geq m. \end{cases}\quad v\in \mathbb{R}^n.\] The existence of the family \(\left\{(\cdot)^{m,\tau}\right\}\) is clear. The \(m\)-truncated, \(\tau\)-smoothed ReLU neuron is defined as: \[\sigma^{m,\tau}(x,\theta)=a^{m,\tau}(x^\top w^{m,\tau}+b)_+^{m,\tau}, \qquad x\in \mathbb{R}^{d},\theta=(a,w,b)\in \mathbb{R}^{d+2}.\]

9.2 An Explicit Construction of the Family \(\left\{(u)_+^{m,\tau}\right\}\)↩︎

We construct a family \(\{(u)_+^{m,\tau}\}\) that approximates the univariate ReLU function and satisfies the requirements in previous subsection.

Select a univariate function \(g\colon [1,+\infty)\to \mathbb{R}\) that satisfies

• \(g(u)\in[1,2],\;\forall u\geq 1\),

• there exists \(\delta>0\) such that \(g(u)=u\) when \(u\in [1,1+\delta]\),

• \(g\in C^4\); for all \(u\geq 1\), \(g'(u)\in [0,1], g''(u)\in[-1,0], |g^{(3)}(u)|\leq C, |g^{(4)}(u)|\leq C\).

Such a function \(g\) implements a truncation. The existence of such a function \(g\) is clear. Let \[S_\tau (u)=\frac{\log(1+e^{\tau u})}{\tau}.\] This is the \(\tau\)-scaled softplus function, which is a smooth approximation of \((u)_+=\max\{0,u\}\). Define the tail function \[\phi_{m,\tau}(u)= m^2 g\Big( \frac{S_\tau(u)}{m^2} \Big).\] Let \[(u)_+^{m,\tau}=\mathbb{I}_{u< c_{m,\tau}} S_\tau(u) + \mathbb{I}_{u \ge c_{m,\tau}} \phi_{m,\tau}(u).\] Here \(c_{m,\tau}=S_\tau^{-1}(m^2)\).

As \(S_\tau(c_{m,\tau})=m^2\) and \(S_\tau\) is monotonically increasing, there exists \(\delta'\) such that \[\frac{S_\tau(u)}{m^2}\in[1,1+\delta],\quad \forall u\in[c_{m,\tau},c_{m,\tau}+\delta'].\] Then, \[\phi_{m,\tau}(u)=S_\tau(u),\quad \forall u\in[c_{m,\tau},c_{m,\tau}+\delta'].\] Therefore, \((u)_+^{m,\tau}\) is \(C^4\). Since \(1\leq g\le 2\), we have that \(|\phi_{m,\tau}(u)|\leq 2m^2\). For the first derivative, \[\phi'_{m,\tau}(u)=g'\Big( \frac{S_\tau(u)}{m^2} \Big) S'_\tau(u).\] Since \(0\le g'\le 1\) and \(0\le S_\tau' \le 1\), we have that \(|\phi'_{m,\tau}(u)|\le 1\). For the second derivative, \[\phi''_{m,\tau}(u)=S''_\tau(u)g'\Big( \frac{S_\tau(u)}{m^2} \Big)+ \frac{1}{m^2}g''\Big( \frac{S_\tau(u)}{m^2} \Big)(S'_\tau(u))^2.\] Since \(S''_\tau < 1/S_\tau\), we have that \[S''_\tau(u) \leq \frac{1}{m^2}, \quad \forall u\geq c_{m,\tau}.\] Since \(-1\le g'' \le 0\), we have that \[\phi''_{m,\tau}(u) \le S''_\tau(u)g'\Big( \frac{S_\tau(u)}{m^2} \Big)\leq \frac{1}{m^2}.\] Also, as \(S''_\tau \geq 0\), \[\phi''_{m,\tau}(u) \ge -\frac{1}{m^2}(S'_\tau(u))\geq -\frac{1}{m^2}.\] Since \(g^{(3)},g^{(4)},S_\tau^{(3)},S_\tau^{(4)}\) are bounded, \(\phi_{m,\tau}^{(3)},\phi_{m,\tau}^{(4)}\) are bounded.

Now define \[\phi_m (u) = m^2 g(\frac{u}{m^2}), \quad u\geq m^2,\] and let \[(u)_+^m= \mathbb{I}_{u<m^2}\cdot (u)_+ +\mathbb{I}_{u\geq m^2} \phi_m(u).\] Since \(S_\tau(u)\to (u)_+\) pointwise and \(c_{m,\tau}\to m^2\) as \(\tau\to+\infty\), we have \((u)_+^{m,\tau}\to (u)_+^m\) as \(\tau\to+\infty\). By similar arguments as above, \((u)_+^m\in C^4((0,+\infty))\). Consider any compact set \(K\subset (m^2,+\infty)\). For all large \(\tau\), \(K \subset (c_{m,\tau},+\infty)\). Note that \[S_\tau(u)\to u,\quad S'_\tau(u)\to 1,\quad S''_\tau(u)\to 0\] uniformly on \(K\) as \(\tau\to+\infty\). Hence, \[\phi''_{m,\tau}(u) \to \frac{1}{m^2}g''(\frac{u}{m^2})=\phi_m''(u)\] uniformly on \(K\) as \(\tau\to+\infty\). By above, the family \(\left\{(u)^{m,\tau}_+\right\}_{m,\tau}\) satisfies all the previous requirements.

10 Proof of Theorem 2↩︎

In this section, we present the proof of Theorem 2. In Appendix 10.1, Proposition [prop:tau-limit] controls the Hessian of the solution function in the zero-smoothing limit; In Appendix 10.2, Proposition [prop:qpc] further bounds the limiting Hessian under the Quadratic Positivity Condition; In Appendix 10.3, Proposition [prop:shrink] shows the shrinkage of the cluster set, which in turn determines where the Quadratic Positivity Condition holds. Combining these ingredients, in Appendix 10.4, we prove Theorem 5, thereby proving Theorem 2. Technical lemmas are provided in Appendix 10.5.

10.1 Proof of Proposition [prop:tau-limit]↩︎

In this section, we prove Proposition [prop:tau-limit] which bounds the Hessian of the solution depending on \(m,\beta,\lambda\) and a term that depends on the training data and residuals. Recall that we let \(\alpha=(\beta,\lambda,m)\). We also let \(\hat{\alpha}=(\beta,\lambda,m,\tau)\).

10.2 Proof of Proposition [prop:qpc]↩︎

In this section, we prove Proposition [prop:qpc] which bounds the Hessian \(\|\nabla_x^2 h^\ast_\alpha(x)\|\) under the Quadratic Positivity Condition.

10.3 Proof of Proposition [prop:shrink]↩︎

In this section, we prove Proposition [prop:shrink] which shows the shrinkage of the cluster set.

10.4 Proof of Theorem 5↩︎

In this section, we prove Theorem 5, thereby proving Theorem 2.

10.5 Technical Lemmas↩︎

In this section, we present technical Lemmas. Lemma 3 shows that the stationary measure \(\rho^\ast_{\hat{\alpha}}=\rho^\ast_{\beta,\lambda,m,\tau}\) has Gaussian tail decay and is locally Lipschitz. It also controls the residuals of \(h^\ast_{\hat{\alpha}}\). Lemma 4 controls the second moments of \(\rho^\ast_{\hat{\alpha}}\) and \(\rho^\ast_{\alpha}\). Lemma 5 shows that \(\rho^\ast_{\hat{\alpha}}\) converges pointwise to \(\rho^\ast_{\alpha}\) as \(\tau\to+\infty\) and that \(\rho^\ast_\alpha\) also has Gaussian tail decay. Lemma 6 bounds the residuals of \(h^\ast_\alpha\). Lemma 7 provides estimates on the partition function \(Z_{\alpha}^\ast\), whose definition is given below. All these technical results are adapted and refined from those of [9] to accommodate arbitrary dimensions.

For \(j=1,\ldots,M\), let \[R^{\hat{\alpha}}_j = \frac{1}{M}(h^\ast_{\hat{\alpha}}(x_j) - y_j).\] Let \[\Psi^\ast_{\alpha}(\theta) = \sum_{j=1}^M R^\alpha_j \cdot \sigma^{m}(x_j,\theta)+\frac{\lambda}{2}\|\theta\|^2, \quad \Psi^\ast_{\hat{\alpha}}(\theta) = \sum_{j=1}^M R_j^{\hat{\alpha}} \cdot \sigma^{m,\tau}(x_j,\theta)+\frac{\lambda}{2}\|\theta\|^2.\] Then \[\rho^\ast_{ \alpha}(\theta) = (Z_{\alpha}^\ast)^{-1}\exp(-\beta\Psi^\ast_{\alpha}(\theta)),\qquad \rho_{\hat{\alpha}}^\ast(\theta) = (Z_{\hat{\alpha}}^\ast)^{-1}\exp(-\beta\Psi^\ast_{\hat{\alpha}}(\theta)),\] where \(Z_{\alpha}^\ast=\int \exp(-\beta\Psi^\ast_{\alpha}(\theta)) d\theta\) and \(Z_{\hat{\alpha}}^\ast=\int \exp(-\beta\Psi^\ast_{\hat{\alpha}}(\theta)) d\theta\).

11 Proof of Theorem 3↩︎

In this section, we present the proof of Theorem 3, which shows that the trained input weights and biases align along a few directions.

12 Proof of Theorem 4↩︎

In this section, we present the proof of Theorem 4. We first prove Lemma 2, which characterizes the geometry of the strata. We then prove Proposition [prop:strat], which provides a geometric proof of Theorem 4.

13 Proof of Proposition [prop:toy]↩︎

In this section, we prove Proposition [prop:toy], which identifies the kink hyperplanes for a simple class of data sets.

14 Comparison of the Bounds on Effective Width↩︎

In this section, we compare our bound on the effective width with the bound obtained in [39]: \[K' \leq K'_{\mathrm{bound}}= \max_{k=0,\ldots,M} \binom{M}{k}\Xi(M-k,d+1-k),\] where \(\Xi(n,d)\) is the number of regions in \(\mathbb{R}^d\) divided by an arrangement of \(n\) hyperplanes that pass through the origin. Note that [39] considered networks without input biases and replace each training input \(x_j\) by \(x_j' =(x_j,1)\in \mathbb{R}^{d'}=\mathbb{R}^{d+1}\). Accordingly, the original bound in [39] is \[\max_{k=0,\ldots,M} \binom{M}{k}\Xi(M-k,d'-k).\]

For the training set \(\mathcal{X}=\left\{x_j\right\}_{j=1}^M\subset \mathbb{R}^{d}\), let \(G_{x_j}=\left\{(w,b)\colon w^\top x_j+b=0\right\}\). Note that in the analysis of [39], \(\Xi(M-k,d+1-k)\) is essentially counting the number of regions of \(\cap_{j\in \mathcal{J}_k} G_{x_j}\) divided by \(\left\{G_{x_j}\right\}_{j\notin \mathcal{J}_k}\) for some \(\mathcal{J}_k \subset \left\{1,\ldots,M\right\}\) with \(|\mathcal{J}_k|=k\) (see the proof of Lemma 3.11 in [39]). Below we consider the general case, i.e., \(\mathcal{X}\) are in general position. In this case, [53] gives that \[2\mathcal{P}(\mathcal{X})-1=2\sum_{i=0}^{d}\binom{M-1}{i} -1,\] and that \[K'_{\mathrm{bound}}=\max_{k=0,\ldots,M} \binom{M}{k} \cdot 2\sum_{i=0}^{d-k}\binom{M-k-1}{i}.\] Note that, \[K'_{\mathrm{bound}}>2\sum_{i=0}^{d}\binom{M-1}{i}.\] Thus, \(K'_{\mathrm{bound}}\) exceeds our bound by at least one.

Now assume that \(2d +1 \leq M\). Note that \[\frac{\binom{M-1}{i+1}}{\binom{M-1}{i}}=\frac{M-1-i}{i+1}\geq 1 \Longleftrightarrow 2i\leq M-2,\] which is satisfied when \(i\leq d-1\). Thus, \[2\mathcal{P}(\mathcal{X}) -1 < 2\sum_{i=0}^d \binom{M-1}{i} \leq 2(d+1) \binom{M-1}{d}.\] Meanwhile, \[\begin{align} \max_{k=0,\ldots,M} \binom{M}{k} \cdot 2\sum_{i=0}^{d-k}\binom{M-k-1}{i} &\geq 2\binom{M}{\lfloor d/2\rfloor}\binom{M-\lfloor d/2\rfloor-1}{d-\lfloor d/2\rfloor}\\ &\geq 2 \frac{M}{M-\lfloor d/2\rfloor} \binom{d}{\lfloor d/2\rfloor}\binom{M-1}{d}\\ &\geq 2 \binom{d}{\lfloor d/2\rfloor}\binom{M-1}{d}\\ &\geq 2 \frac{2^d}{d+1}\binom{M-1}{d}. \end{align}\] Therefore, \[\frac{K'_{\mathrm{bound}}}{2\mathcal{P}(\mathcal{X}) -1} \geq \frac{2^d}{(d+1)^2}.\] That is, in the regime of \(2d +1 \leq M\), our bound improves over the bound in [39] by a factor of at least \(2^d/(d+1)^2\).

15 Existence of Accumulation Points↩︎

In this section, we show that for any sequence \((\beta_n,\lambda_n,m_n)\) satisfying Assumption 1, with \(m_n,\beta_n\to+\infty\) and \(\lambda_n\to \bar \lambda\), one can find a subsequence \((\beta_{n_k},\lambda_{n_k},m_{n_k})\) along which \(h^\ast_{\beta_{n_k},\lambda_{n_k},m_{n_k}}(\cdot)\) converges pointwise.

16 Non-redundancy of Kink Hyperplane Arrangement↩︎

In this section, we obtain a description of the location of the kink hyperplanes of the limiting piecewise affine function using Theorem 4.

We characterize the location of a hyperplane in the input space through the partition it induces on the training set \(\mathcal{X}=\left\{x_j\right\}_{j=1}^M\). As defined in Section 4.1, a dichotomy is a partition of the training set induced by a hyperplane that does not pass through any data point. To accommodate hyperplanes that intersect the training set, we introduce the following notion. For a training set \(\mathcal{X}=\left\{x_j\right\}_{j=1}^M\), we say \(\mathcal{T}=(\left\{\mathcal{X}_+,\mathcal{X}_-\right\},\mathcal{X}_0)\) is a linear trichotomy of \(\mathcal{X}\) if \(\mathcal{X}=\mathcal{X}_+\sqcup \mathcal{X}_-\sqcup \mathcal{X}_0\) and there exist \(w\in\mathbb{R}^d,b\in \mathbb{R}\) satisfying \(w^\top x_j +b>0\) if \(x_j \in \mathcal{X}_+\), \(w^\top x_j +b<0\) if \(x_j \in \mathcal{X}_-\), and \(w^\top x_j +b=0\) if \(x_j \in \mathcal{X}_0\). In this case, we say that \(\mathcal{T}\) is realized by the hyperplane \(\{x\colon w^\top x+b=0\}\). Note that a dichotomy is a special trichotomy with \(\mathcal{X}_0 = \varnothing\). We call a trichotomy a strict trichotomy if it is not a dichotomy. We introduce a partial order to compare trichotomies. For two trichotomies of \(\mathcal{X}\), \(\mathcal{T}=(\left\{\mathcal{X}_+,\mathcal{X}_-\right\},\mathcal{X}_0)\) and \(\mathcal{T}'=(\left\{\mathcal{X}_+',\mathcal{X}_-'\right\},\mathcal{X}_0')\), define \(\mathcal{T}'\preceq \mathcal{T}\) if, after possibly interchanging \(\mathcal{X}_+\) and \(\mathcal{X}_-\), \[\mathcal{X}_+\subseteq \mathcal{X}_+',\quad \mathcal{X}_-\subseteq \mathcal{X}_-'.\] When \(\mathcal{T}'\preceq \mathcal{T}\), we say that \(\mathcal{T}\) is a degeneration of \(\mathcal{T}'\).

Recall from Section 6.5 that, each stratum \(V\in \mathcal{V}\) uniquely corresponds to a realizable ternary sign pattern on the training inputs \(\mathcal{X}\). Note that each realizable trichotomy of \(\mathcal{X}\) corresponds to a pair of ternary sign patterns related by a sign flip, and hence to a pair of antipodal strata \(\{V,-V\}\), except that the trichotomy \((\{\varnothing,\varnothing\},\mathcal{X})\) corresponds to the single stratum \(V_0\). Under this correspondence, the following result is a direct consequence of Theorem 4.

[9] showed that the solution function has at most three knots between any two consecutive training points and, when three knots occur, two must coincide with the endpoints of the interval, while the third lies strictly in its interior. We explain that this is covered as a special case of Corollary 1. Consider a set of univariate training inputs \(\mathcal{X}=\left\{x_j\right\}_{j=1}^M\) and assume that \(-L=x_0<x_1<\cdots<x_M<x_{M+1}=L\) for some sufficiently large \(L>0\). Fix a \(j\in \left\{1,\ldots, M-1\right\}\). Let \(\mathcal{T}\) denote the dichotomy \(\mathcal{X}=\left\{x_{j'}\right\}_{j'\leq j}\sqcup \left\{x_{j'}\right\}_{j'>j}\). Note that \(|\mathcal{D}(\mathcal{T})|\) precisely counts the knots in the closed interval \([x_{j},x_{j+1}]\). By Corollary 1, if \(|\mathcal{D}(\mathcal{T})|=3\), then at least two of the corresponding hyperplanes realize strict trichotomies. In the univariate setting, this means that two of the associated knots must lie exactly at \(x_{j}\) and \(x_{j+1}\), while the remaining knot must lie in \((x_{j},x_{j+1})\). On the other hand, if \(|\mathcal{D}(\mathcal{T})|=4\), then Corollary 1 implies that all four corresponding hyperplanes must realize strict trichotomies. This is impossible on the interval \([x_{j},x_{j+1}]\), since only the end points correspond to strict trichotomies. Regarding the dichotomy \(\mathcal{T}_0=(\left\{\mathcal{X}, \varnothing\right\},\varnothing)\), note that \(\mathcal{D}(\mathcal{T}_0)\) counts the knots in the intervals \([x_0, x_1]\) and \([x_M,x_{M+1}]\). By Corollary 1, there are at most two knots on these two closed intervals. Moreover, if two knots are present, they must occur precisely at the boundary data points \(x_1\) and \(x_M\).

References↩︎