Scalable optimization for large data

[stli]

First a big thanks to all maintainers for this great library. I am trying to optimize a cost function on potentially larger scale input data for cases where classical mini-batch optimizers converge very slowly. For mini-batch optimizers (e.g. optax), this could be easily done using a dataloader (e.g. grain). However, it seems a bit tricky to do in optimistix. Hoping for some insights and/or ideas how to achieve this with optimistix.

Below is an MWE and a list of my considered variants

1. gpu sharding (needs more physical gpus)
2. move to cpu (plenty of mem but slow)
3. alternating solves for different data batches (needs some logic to ensure convergence)
4. alternating solves for different nn layers (needs some logic to ensure convergence)
5. dataloader to calculate cost function iteratively (seems ideal but difficult to integrate)

import jax
import jax.numpy as jnp
import optimistix as optx
import jax.random as jr
from jax.example_libraries import stax

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

key = jr.PRNGKey(0)
dim = 10
num_samples = 100_000 # some large number

X = jr.normal(key, (num_samples, dim))
y = jr.normal(key, (num_samples, 1))

activation = stax.Sigmoid
layer_sizes = [80,]*3
init_fun, predict_fun = stax.serial(
    stax.serial( *(sum([[stax.Dense(size), activation] for size in layer_sizes],[],)) + [stax.Dense(1)] ) )
_, params = init_fun(key, (X.shape[1],))

def loss(params, args):
    X, y = args # <- ideally an iterator + fori loop over data batches here
    return jnp.mean( (predict_fun(params, X).flatten() - y.flatten())**2 )

bfgs_tol = 1e-12
solver = optx.BFGS(rtol=bfgs_tol, atol=bfgs_tol)
sol = optx.minimise(
    loss,
    solver,
    params,
    max_steps=100_000,
    throw=False,
    args=(X,y) # <- ideally a dataloader here
)

[patrick-kidger]

When your cost function is given by an average loss over a large dataset, then this is often framed as a stochastic/noisy optmization problem – as on each step the loss is approximated with a Monte-Carlo estimate.

In this day and age, such problems are pretty much ubiquitously solved using first-order gradient methods, as already provided by Optax. (Meanwhile the optimizers of Optimistix generally aren't written for such scenarios, and tend to focus on other problems, e.g. for which second order methods are appropriate.)

For this reason I think I'd most recommend just sticking to Optax for such problems. I don't think there's much we can do to help you beyond what they already provide :)

[stli]

Many thanks for your response. I agree that stochastic first-order methods are indeed the best choice for many contemporary large-scale problems. One notable exception seems to be PINNs or other learning tasks involving differential operators applied to neural networks. In these cases, second-order methods such as (L-)BFGS and its variants seem still preferred over contemporary first-order methods (see, e.g., the paper “Challenges in Training PINNs: A Loss Landscape Perspective” by P. Rathore).

I understand this is not the main focus of optimistix. However, I am curious if there is any advice on how one might reduce peak memory usage with optimistix solvers? To be more specific, do you think a custom vmap, such as jax.custom_batching.sequential_vmap or an adapted variant, could work?

[patrick-kidger]

Okay! So if your goal is simply to reduce peak memory usage, by computing batch elements sequentially rather than in parallel, then I agree that something like a custom-batching rule (or a manual lax.scan inside the f that we optimize) would be one approach.

I think it'd be important that this iterate over the whole dataset on every step, though – I'm not sure how well these second-order optimizers work if they're given noisy estimates at every step.