Mini Batch with OptaxSolver

[jecampagne]

Hello,
Let me give you a snippet

import jax
import jax.numpy as jnp
import jaxopt

from jax.config import config
config.update("jax_enable_x64", True)

import numpy as np
#######
def ridgeless_reg_objective(params, X, y):
  residuals = jnp.dot(X, params) - y
  return jnp.mean(residuals ** 2)

def gen_x_unif_sphere(r=1.0, d=20,ns=50, seed=42):
    """ Generate ns vectors uniform on the sphere of radius r in d-dimension
        d=20.
        data = gen_x_unif_sphere(r=np.sqrt(d), d=d)
        jnp.allclose(jnp.sum(data*data,axis=1),jnp.array([d]*data.shape[0]))
        
    """
    key = jax.random.PRNGKey(seed)
    x = jax.random.normal(key, shape=(ns,d))
    norm = jnp.linalg.norm(x, axis=1)
    x_normed = r * x / norm.reshape(x.shape[0],1)
    return x_normed

def func(x):
    def p1(x):
        return x
    def p2(x):
        return (x*x - 1.)/jnp.sqrt(2.)
    def p3(x):
        return x*(x*x-3.)/jnp.sqrt(6.)
    def p4(x):
        return (x*x*x*x - 6. * x*x +3.)/(2.*jnp.sqrt(6.))
    c1 = jnp.sqrt(2./5.)
    c2 = jnp.sqrt(1./5.)
    return c1*(p1(x)+p2(x))+c2*p3(x)


def gen_y(X, beta):
    Xbeta = X @ beta.T  # <beta, Xi>
    y = func(Xbeta)
    return y

def rho_prime(x):
    """
        rho(x)=ReLU(x) => rho'(x)=Heaviside(x) with  by convention ReLU'(0)=0.
    """
    return jnp.heaviside(x, 0.)

def gen_Phi(X, W):
    N, d = W.shape
    Nd = N*d
    ns = X.shape[0]
    XW = X @ W.T
    rhoXW = rho_prime(XW)/jnp.sqrt(Nd)
    return jnp.tile(X,reps=N) * jnp.repeat(rhoXW,repeats=d, axis=1)

##########
d = 15
# gamma = ln(n)/ln(d)
# psi   = ln(Nd)/ln(d)
gamma = 3.0
ns = int(d**gamma)
psi = 4.0
N = int(d**(psi-1))
print(f"gamma:{gamma}, d:{d}, ns:{ns}, N:{N}, Nd:{N*d}")

X = gen_x_unif_sphere(r=np.sqrt(d), d=d, ns=ns)
print("X.shape",X.shape)

key = jax.random.PRNGKey(70)
beta = jax.random.normal(key, shape=(1,d))
norm = jnp.linalg.norm(beta, axis=1)
beta =  beta / norm


Y = gen_y(X, beta)
print("Y.shape",Y.shape)

W = gen_x_unif_sphere(r=1,d=d,ns=N,seed=60)
print("W.shape",W.shape)


Phi= gen_Phi(X,W)
print("Phi.shape",Phi.shape)

#######

gives

gamma:3.0, d:15, ns:3375, N:3375, Nd:50625
X.shape (3375, 15)
Y.shape (3375, 1)
W.shape (3375, 15)
Phi.shape (3375, 50625)

Now make a regression Y = Phi Theta using Adam

import optax
from jaxopt import OptaxSolver
init_theta = jnp.zeros_like(theta)
solver = OptaxSolver(opt=optax.adam(1e-2), fun=ridgeless_reg_objective,maxiter=1000)
theta_adam = solver.run(init_theta, X=Phi, y=Y)

Of course you can try a least squared min-norm to get theta but even here the matrix are already sizeable. Now, with the LSQ solution I get a train MSE error computed

diff_lsq = Phi @ theta_lsq - Y
MSE_lsq = (diff_lsq.T @ diff_lsq)/diff_lsq.shape[0]

which is equal to `DeviceArray([[5.98907097e-28]]l and is fine as Nd>>ns.

While with Adam I get

diff_adam = Phi @ theta_adam.params - Y
MSE_adam = (diff_adam.T @ diff_adam)/diff_adam.shape[0]

leads to DeviceArray([[0.01393647]]) which is far worse.

I wander if there is some batch options or some parameters to tune to get better results. Thanks.