AI Training as Nonconvex Optimization

Train two hidden-layer parameters of a small neural network

Forty-eight regression samples are generated by a four-unit softplus teacher network with mild deterministic noise. A narrower three-unit student network is trained by empirical squared loss with weight decay. Only two hidden-layer weights move, giving an exact two-dimensional nonconvex training problem.

Trainable coordinatesu = input weight of hidden unit 1
v = input weight of hidden unit 2

Activationsoftplus(z) = log(1 + exp(z))
smooth, so Newton has a genuine Hessian

ŷ(x;u,v)=0.1+1.2·sp(ux−0.1)−0.8·sp(vx+0.4)−0.3·sp(1.2x+0.15)

L(u,v)= (1/2n) Σᵢ(ŷ(xᵢ;u,v)−yᵢ)² + 0.015(u²+v²)/2

× selected initial pointcolored paths = optimization iterates

Training-loss contoursClick to choose (u₀,v₀)

Convergencetraining loss L(θₖ)

Outcome after the selected iterations

Method	L	‖∇L‖₂	‖θ−θ★‖₂

Step size α0.153

Momentum0.85

Iterations100

SGD and Adam use reproducible mini-batches of 8. The gradient-norm view evaluates the full training gradient ∇L at every iterate. Newton uses α as damping. Cubic Newton uses α=1/L, where L is the cubic-regularization threshold in (L/6)‖p‖³.