Optimization, Generalization and Differential Privacy Bounds for Gradient Descent on Kolmogorov-Arnold Networks

AI summary

Overview

Research area: Machine learning theory — specifically the training dynamics, statistical generalization, and differential privacy of Kolmogorov–Arnold Networks (KANs), a structured alternative to standard multilayer perceptrons (MLPs).

Technical level: Advanced. The paper is a theoretical learning-theory contribution relying on NTK (neural tangent kernel) analysis, Hessian/curvature control, algorithmic stability, and differential privacy accounting. The core claims are stated as theorems with explicit width, step-size, and iteration scalings.

One-sentence scope: The paper derives optimization, generalization, and differential-privacy utility bounds for gradient descent (and its private variant) on two-layer KANs with B-spline edge functions, and shows that polylogarithmic width suffices — and, under privacy, is essentially necessary.

What This Paper Is About

KANs replace the fixed pointwise nonlinearities of MLPs with learnable univariate functions on network edges, which gives them a distinct structural bias and strong empirical results in structured domains. However, most existing KAN theory is algorithm-independent, addressing expressivity and approximation rather than what gradient descent actually produces. This paper closes that gap: it analyzes gradient descent (GD) and differentially private gradient descent (DP-GD) for training two-layer KANs, establishing how network width, iteration count, and privacy budget interact with optimization error and population risk.

Key Contributions

1. A unified reference-point framework for GD and DP-GD on KANs. The authors analyze training relative to a reference point Θ*, yielding general bounds on the optimization risk (training loss), the population risk, and the trajectory-averaged population risk under (ε, δ)-DP (Theorems 4.3, 4.11, and 4.15).

2. Optimization and fast-rate generalization under NTK separability with logistic loss. They show that polylogarithmic width suffices for GD to achieve an optimization risk of order Õ(1/T) (Theorem 4.9) and a fast-rate population risk bound of order Õ(1/n) (Theorem 4.12). The paper states this is the first algorithm-dependent generalization bound specifically for GD-trained KANs.

3. Differential privacy with an explicit noise calibration and utility tradeoff. They provide a Gaussian noise calibration ensuring (ε, δ)-DP for DP-GD on KANs, and derive a utility bound of order Õ(√d/(nε)) for the trajectory-averaged population risk, which matches the classical lower bound for general convex Lipschitz problems (Theorem 4.16). To the authors' knowledge, this is the first utility analysis of DP gradient-based methods for KANs.

4. A sharp width characterization. In the NTK-separable setting, the authors show polylogarithmic width is not only sufficient but essentially necessary for DP-GD to reach the desired utility — a qualitative gap between non-private GD (sufficiency only) and DP-GD (necessity also emerges).

5. Theory-guided practical implications, validated empirically. The theory predicts diminishing returns from increasing width m and iterations T beyond moderate values for GD, whereas under a fixed privacy budget DP-GD benefits from moderate widths and early stopping. These trends are illustrated in the experiments (Section 5).

Main Findings

• Model class: The analysis covers a two-layer KAN of width m with scalar output, input dimension d, and edge functions parameterized by a B-spline basis of size p, with a bounded activation σ (e.g., tanh or sigmoid) after the first layer. The total parameter vector lies in R^{mp(d+1)}. An explicit 1/√d normalization is added in the first layer to stabilize pre-activations as d grows, and a 1/√m factor normalizes the output.

• Optimization rate: Under NTK separability with margin γ > 0, polylogarithmic width suffices for GD with constant step size η to attain an optimization risk of order Õ(1/(γ²ηT)); the informal statement is Õ(1/(γ²ηT)) and the precise version is ≲ (log²(T) + log(n/δ))/(γ²ηT).

• Required width for optimization: m ≳ log(m/δ)(log⁶(T) + log³(n/δ))/γ⁶ — a polylogarithmic dependence on n and T.

• Generalization rate: With ηT ≳ n and the same width scaling, the expected population risk of the GD-trained KAN is ≲ (log⁴(n) + log²(n/δ))/(γ⁴n), i.e., Õ(1/n) up to logarithmic factors.

• Implicit regularization: The generalization bound is stated in terms of the reference-point complexity, so the analysis favors predictors with low population risk that remain close to initialization. This is described as an implicit regularization effect of GD.

• Generalization gap mechanism: The gap L(Θ(T)) − L_S(Θ(T)) is bounded via on-average argument stability, and is controlled by the cumulative training loss along the GD trajectory — so whenever GD achieves small training loss, it also has a small generalization gap.

• Stay-in-a-ball property: Theorem 4.3 additionally shows that GD iterates remain within a controlled neighborhood of both the initialization Θ(0) and the reference point Θ* (distances bounded by √2·Λ_Θ* and 3·Λ_Θ* respectively). This property is described as the cornerstone of the analysis, enabling local smoothness and almost-convexity along the trajectory.

• Trajectory-dependent curvature: Explicit Hessian calculations show both the minimum and maximum eigenvalues of the training loss depend on the evolving deviation ‖Θ(k) − Θ(0)‖₂. The authors handle this with a double induction that simultaneously bounds the cumulative training loss and the distances to Θ* and Θ(0).

• Realizability variant: Under a realizability assumption parameterized by g(ε) = inf{‖Θ − Θ(0)‖₂ : L_S(Θ) ≤ ε}, the optimization risk is bounded by (η + g²(1/T))/(ηT). If g(ε) ≲ log(1/ε), then ηT ≍ n yields L_S(Θ(T)) ≲ log²(n)/n with m ≳ log⁶(n).

• Privacy mechanism: DP-GD is implemented by adding Gaussian noise to the gradients each iteration and projecting the updates onto bounded domains Ω and Ω_c, rather than using gradient clipping; sensitivity is controlled by constraining the iterates via projection. The noise variance takes the form σ̃² = T(1 + log(2T/δ)/ε)(n²ε)^{-1}, with per-coordinate scaling constants C₁ and C₂.

• Private utility bound: For ηT ≍ γ²nε/√d and polylogarithmic width m ≍ polylog(n), the (ε, δ)-DP variant achieves an expected population risk averaged over T iterates bounded by Õ(√d/(γ⁴nε)).

• Comparison to the classical lower bound: The paper notes that a minimax lower bound tailored to KANs under DP is not currently available, so the private utility bound is compared against the classical lower bound for general convex Lipschitz problems (Bassily et al., 2019); the KAN bound matches it up to the separability margin and logarithmic factors.

• Contrast with prior KAN optimization theory: Gao and Tan (2025) proved global linear convergence (1 − (η/2)λ_min(G^∞))^T for two-layer B-spline KANs in regression, but under a positive-definiteness assumption on the expected NTK Gram matrix requiring widths that scale polynomially in problem parameters. This paper instead uses NTK separability, described as weaker than Gram-matrix positive-definiteness, and achieves sublinear Õ(1/(ηTγ²)) convergence already at polylogarithmic width.

• Empirical illustration: Figure 1 compares MLP and KAN performance on genomic sequence classification across benchmark datasets from Cherednichenko and Poptsova (2025), using the Matthews correlation coefficient (MCC, higher is better). Most KAN points lie to the right of their MLP counterparts, indicating improved predictive performance when MLP modules are replaced with KAN layers. The specific numerical experiment results beyond this figure are not reported in the truncated content.

Methodology in Plain English

The authors study a concrete, analyzable KAN: two layers, one hidden layer of width m, scalar output, and each edge function written as a B-spline expansion. Everything is collected into a single flat parameter vector, and training is plain gradient descent with a constant step size from a standard Gaussian initialization.

The central technical difficulty is that the loss surface the iterates see is not fixed — its curvature depends on how far the parameters have moved from where they started. The authors attack this with a "double induction": they simultaneously prove that (i) the accumulated training loss stays small and (ii) the parameters stay inside a bounded ball around initialization and around a chosen reference solution. Once both hold together, the Hessian eigenvalues are uniformly controlled along the trajectory, so the loss behaves smoothly and almost convexly along the path that GD actually takes, and a standard convergence argument goes through.

To turn this into concrete rates, they assume NTK separability — a margin condition saying there is a direction in the gradient feature space at initialization that separates the classes by at least γ. This is a weaker requirement than the usual assumption that the NTK Gram matrix is positive definite. Under this assumption, the reference-point distance scales polylogarithmically, and the rates become 1/T for optimization and 1/n for population risk, with width that grows only polylogarithmically in the sample size and iteration count.

For generalization, they compare the behavior of the algorithm on a dataset and on a dataset with one point changed (on-average argument stability), showing the generalization gap is controlled by the cumulative training loss. For privacy, they add carefully calibrated Gaussian noise at every step and clip by projection onto fixed balls instead of gradient clipping, then track how the noise propagates through the nonconvex trajectory. That propagation controls both the privacy guarantee and the averaged population risk.

Why This Matters

Impact on research. The paper supplies what the authors describe as the first algorithm-dependent generalization bounds for GD-trained KANs and the first utility analysis of DP gradient methods for KANs. It also identifies a structural asymmetry: polylogarithmic width is merely sufficient in the non-private regime, but essentially necessary once differential privacy is imposed, which suggests that privacy changes the required model capacity, not just the obtained accuracy.

Real-world applications (as motivated by the paper):

• Genomic sequence classification, where the paper's Figure 1 shows KAN layers improving MCC over MLP modules on benchmark datasets from Cherednichenko and Poptsova (2025).
• Scientific computing and physics-informed learning, where interactions can be represented through univariate components.
• Time series forecasting via temporal KAN designs.
• Molecular and biological modeling, plus vision tasks using convolutional KAN layers and Transformer architectures whose feedforward subnetwork is replaced by KAN layers.

Industry relevance. Since the motivating domains include biology and medicine, where training data are often sensitive, the differential privacy analysis is directly relevant to deploying KAN-based models on private data. The width-selection and early-stopping guidance — diminishing returns from larger width and more iterations for GD, and moderate width plus early stopping under a fixed privacy budget — gives practitioners concrete tuning rules grounded in theory rather than trial and error. Differential privacy provides the rigorous framework used here for limiting the influence of any single individual's data point.

Future Directions

• A KAN-specific DP minimax lower bound. The paper explicitly notes that a minimax lower bound tailored to KANs under DP is not currently available, so the private utility bound is compared to the classical convex Lipschitz lower bound of Bassily et al. (2019) instead. Establishing a KAN-specific lower bound is a natural open problem.

• Extending beyond two layers. The analysis is restricted to two-layer KANs with B-spline edge functions and a scalar output. Whether the polylogarithmic-width story, the stay-in-a-ball property, and the privacy necessity gap carry over to deeper KANs is open.

• Sharpening the necessity statement. The necessity result is described as holding "essentially" and in the NTK-separable setting with an admissible width regime and an essentially matching lower bound; tightening this characterization and testing whether it holds beyond NTK separability would strengthen the qualitative claim.

• Connecting the theory to the experiments. The truncated content states that experiments in Section 5 support the qualitative trends for width m and iteration number T, but the quantitative details and dataset sizes are not reported in the available text; a fuller empirical validation of the predicted width and stopping rules is an evident next step.

Target Audience

This paper is aimed at learning-theory researchers working on neural network optimization, generalization, and differential privacy, particularly those interested in the theoretical foundations of KANs. It is also relevant to researchers in differentially private machine learning who need utility guarantees for gradient-based training of non-MLP architectures, and to advanced practitioners in scientific computing, genomics, and medical modeling who want principled guidance on choosing KAN width and stopping time — especially under a privacy budget. Readers need comfort with NTK analysis, convergence rates, and DP definitions to follow the theorems, though the overview section (Section 3) and the informal theorem statements make the high-level messages accessible to a broader machine learning audience.

Authors’ abstract

Kolmogorov--Arnold Networks (KANs) have recently emerged as a structured alternative to standard MLPs, yet a principled theory for their training dynamics, generalization, and privacy properties remains limited. In this paper, we analyze gradient descent (GD) for training two-layer KANs and derive general bounds that characterize their training dynamics, generalization, and utility under differential privacy (DP). As a concrete instantiation, we specialize our analysis to logistic loss under an NTK-separable assumption, where we show that polylogarithmic network width suffices for GD to achieve an optimization rate of order $1/T$ and a generalization rate of order $1/n$, with $T$ denoting the number of GD iterations and $n$ the sample size. In the private setting, we characterize the noise required for $(ε,δ)$-DP and obtain a utility bound of order $\sqrt{d}/(nε)$ (with $d$ the input dimension), matching the classical lower bound for general convex Lipschitz problems. Our results imply that polylogarithmic width is not only sufficient but also necessary under differential privacy, revealing a qualitative gap between non-private (sufficiency only) and private (necessity also emerges) training regimes. Experiments further illustrate how these theoretical insights can guide practical choices, including network width selection and early stopping.