Not working: Minimizing a loss function using Levenberg–Marquardt algorithm

Hi,

I am attempting to train a neural network to approximate a 3D function in which I trying to minimize the loss function using the Levenberg-Marquardt (LM) optimizer in Optimistix. However, the parameters do not update beyond the first optimization step. The loss remains unchanged after the first iteration, indicating that the optimizer may not be properly updating the parameters. You can find my code below. Not sure where the error lies.

import jax.numpy as jnp
import jax.random as random
import optimistix
import numpy as np
import matplotlib.pyplot as plt
from flax import linen as nn
from sklearn.preprocessing import MinMaxScaler

Generate Synthetic 3D Data

key = random.PRNGKey(42)
X_data = jnp.linspace(-5, 5, 20)
Y_data = jnp.linspace(-5, 5, 20)
Z_data = jnp.linspace(-5, 5, 20)

X_mesh, Y_mesh, Z_mesh = jnp.meshgrid(X_data, Y_data, Z_data, indexing='ij')
X_flat, Y_flat, Z_flat = X_mesh.flatten(), Y_mesh.flatten(), Z_mesh.flatten()

def true_function(X, Y, Z):
return 2.0 * jnp.sin(1.5 * X) + 1.0 * jnp.cos(2.0 * Y) + 0.5 * Z

F_values = true_function(X_flat, Y_flat, Z_flat) + 0.05 * random.normal(key, (len(X_flat),))
X_train = jnp.stack([X_flat, Y_flat, Z_flat], axis=1)
F_train = jnp.array(F_values)

scaler = MinMaxScaler()
X_train = jnp.array(scaler.fit_transform(np.array(X_train)), dtype=jnp.float32)

Define Neural Network Model

class TanhMLP(nn.Module):
layers: list

@nn.compact
def __call__(self, x):
    for units in self.layers[:-1]:
        x = nn.Dense(units)(x)
        x = jnp.tanh(x)
    return nn.Dense(self.layers[-1])(x)

key = random.PRNGKey(42)
model = TanhMLP(layers=[3, 128, 128, 64, 1])
params = model.init(key, jnp.ones((1, 3), dtype=jnp.float32))

Define Residual Function

def residuals(params, args):
X, true_F = args
predicted_F = model.apply(params, X).flatten()
return [predicted_F - true_F]

Optimize Using Levenberg-Marquardt

solver = optimistix.LevenbergMarquardt(rtol=1e-4, atol=1e-4, verbose=frozenset(["loss", "step"]))
solution = optimistix.least_squares(residuals, solver, params, args=(X_train, F_train), max_steps=2000)
optimized_params = solution.value

It would be great if I could get some help here!

Start by running the provided reproducer around optimistix.LevenbergMarquardt and optimistix.least_squares with the shown neural-network residuals. Trace the first optimization step and parameter value, then determine why later steps do not change the loss. Done means the cause is identified and a regression check demonstrates that parameters and loss continue to update.