The Proof of Kolmogorov-Arnold May Illuminate NN Learning

The Kolmogorov-Arnold theorem (KA) [kol56, arn57] resolved Hilbert’s 13th problem in the context of continuous functions (Hilbert’s literal statement). It took some time for applied mathematicians and computer scientists to notice that it addresses the representation power of shallow, but highly non-linear, neural nets [hed71]. The relevance of KA to machine learning (ML) has been much debated [gp89, v91].

The primary criticism is that the activation functions are not even first differentiable (even when the function $f$ to be represented is); indeed, they must be quite wild, and appear impossible to train. I agree with this criticism, but will argue here that it misses a more important point. The discussion of KA relevance has revolved around its statement. But in mathematics, proofs are generally more revealing of power than statements; I believe this particularly true in the case of KA. The purpose of this note is to call attention to some wisdom embedded within the proof which I believe will be useful for the training of neural nets (NNs). As far as resurrecting KA from the dustbin of ML history, this has already been done earlier this year in [ZWV24], where the original 1-hidden layer of KA has been deepened to many layers while the interlayer maps somewhat tamed (the authors use cubic $b$-splines).

Before explaining KA’s insight and its potential applicability, let me give a modern, optimized statement of KA quoted from [mor21]. This statement is perhaps unnecessarily succinct from the NN perspective: it gets by with a single outer function $g$, not $2n+1$ outer functions which might appear more natural. Also, in the NN context one would expect the inner functions to carry a second index and not merely be rationally independent scalars times a single index function $\phi_{k}$. But we take this statement as representative and refer the reader to Morris for more historical information on the individual contributions.

The proof of KA is divided into the construction of inner functions $\phi_{i}$ and an outer function $g$. The inner functions taken together constitute an embedding, which we now denote $\Phi:\mathbb{I}^{n}\rightarrow\mathbb{R}^{m}$, $m\geq 2n+1$. We will append a subscript $\Phi_{i}$ when referring to a neuron coordinate of the embedding. $\Phi$ is universal; it can be chosen once and for all, independent of the function $f$ to be represented. Our chief lesson regards $\Phi$, although I also will make a comment on the non-linear dynamic used to converge to $g$ (given $f$).

In the mathematically idealized case of KA, feasible $\Phi$ are actually dense in the space of continuous maps $\mathcal{C}(\mathbb{I}^{n},\mathbb{R}^{m})$; they also have a crucial local feature, about to be described. This fact is reminiscent of the observation (for example see [gmrm23, skmf24]) that when NNs are randomly initialized, training is often seen to merely perturb the values of the early layers. These layers seem to be more in the business of assuming a form conducive to the learning of later layers, rather than attending to the particular data themselves.

The local property crucial to $\Phi$ is a kind of irregular staircase structure as might be used to approximate the general continuous function by one which is piecewise constant. More specifically, $\Phi$ may be taken to be Lipschitz [ak17], and so by Rademacher’s theorem will be differentiable a.e. Such $\Phi$ constructed in the proof of KA will have the property that at every point $x\in\mathbb{I}^{n}$ where the Jacobian is defined it will, as an $n\times m$ matrix in the “neuron basis,” have $m-n$ columns consisting entirely of zeros. So only one of its ($m$ choose $n$) $n\times n$ minors can be nonzero. (Of course, which minor is active will vary for with $x$.) To picture what this means, imagine an irregular stack of sugar cubes of different sizes but with all faces perpendicular to one of the three coordinate axes, $x$, $y$, or $z$. Another condition says that the “inactive” coordinate on the $w$-perpendicular faces ($w=x$, $y$, or $z$) must take distinct values. The actual situation is more like an irregular structure of $n$-D cubes wafting through $m$-D space, now with $m-n$ inactive directions at any generic point, and all these inactive coordinate values distinct. Being distinct is what sets the stage towards finding at least an approximate outer function $g_{0}$. The $m-n$ inactive directions are crucial to the construction of $g_{0}$, and it is important that $(m-n)>n$, so that the inactive coordinates are in the majority at all generic points of $\Phi(x)$.

The very rough idea for how to choose $g_{0}$ is to use the $(m-n)$ inactive function value on each “sugar cube face” as a hint: If that value $f(x)$ is uniformly positive (negative) on the face with inactive coordinate $\Phi_{i}(x)$, give $g_{0}(\Phi_{i}(x))$ a small constant positive (negative) value and extend in a convex-PL manner. In the regime $(m-n)>n$, it may be checked that such a strategy produces a $g_{0}$ leading to a useful approximation:

Now iterating on the approximation error (in a manner reminiscence of Resnets), and relying on additivity at the output node, one produces a series $g_{0},g_{1},\dots$ which converges exponentially fast to the desired outer function $g$ for which the approximation error has been driven to zero. Many inactive directions at each point are essential since they provide a stationary coordinate value on which a guess for $g_{0}$ can be made: A small positive constant if $f(x_{1},\dots,x_{n})$ is positive throughout the face and a small negative constant if $f$ is negative throughout the face.

This sketch of the construction of the outer function g reveals the importance of the stationarity condition and its implication of vanishing Jacobian minors in the embedding $\Phi$. Remembering that admissible $\Phi$ are dense, we see that the job of the first layer is simply to impress a certain microscopic texture on $\Phi$ which makes g constructable (if not learnable). This is the division of labor that feels so striking. $\Phi$ does not try to learn anything but sets up for success the task of finding $g$.

I wonder if a similar division of labor occurs between early and later stages of a NN as it is trained, and propose:

First, what are these natural Jacobians, and second, what is minor concentration? One may consider either a conventional feedforward DNN or the recently proposed KANets [ZWV24], or more elaborate architectures with residual connections and self-attention. The idea is to watch the data manifold flow through the NN. Initially at $t=0$, the data manifold is $\mathbb{I}^{n}$, the values stored in the $n$ input neurons, and at time $j$ when the input has been imaged in the $j$th layer (say after application of the activation functions). Let $k_{ij}$, $0\leq i,j<\ell_{\text{max}}$, be the map between layers $i$ and $j$, with $i=0$ being a very interesting special case. At any point $y$ in the evolving data manifold ($y=x$, the previously defined coordinate on $\mathbb{I}^{n}$, if $i=0$) we may consider $J_{y}(k_{ij}):=J_{ij}y$, the Jacobian matrix (in neuron coordinates) of the forward map $k_{ij}$ at point $y$. The proposal is to find some context where it is possible to search (or at least sample) this huge family of Jacobians and look for minor concentration. Based on understanding KA one may expect, particularly in some initial segment of the NN (perhaps all the way up to the penultimate layer), that training will cause these Jacobians to have their minors concentrated far beyond what would be seen in a random collection (e.g. Gaussian distributed entries) of linear maps. It seems best at first to start from the beginning and set $i=0$. That is, one should look “in the wild” to see if a KA-style of data preparation is discovered, at least for certain tasks, and certain architectures, as a result of training.