Choose optimizers

liesel.optim.LieselOptim requires an explicit optimizer. Pass a configured Optax transformation, such as optimizers=optax.adam(0.01), to use it for all parameters, or optimizers="lbfgs" for a full-data, deterministic fit. L-BFGS cannot use minibatches.

This page uses a small Gaussian regression. Expand the setup to run the examples from top to bottom.

import logging

import jax.numpy as jnp
import numpy as np
import optax
import tensorflow_probability.substrates.jax.distributions as tfd

import liesel.model as lsl
import liesel.optim as opt

Show code cell source

Hide code cell source

logging.getLogger("liesel").setLevel(logging.WARNING)

rng = np.random.default_rng(42)
x = np.linspace(-1.0, 1.0, 128)
X = lsl.Var.new_obs(jnp.asarray(x), name="X")
beta = lsl.Var.new_param(0.0, lsl.Dist(tfd.Normal, 0.0, 5.0), name="beta")
log_sigma = lsl.Var.new_param(0.0, name="log_sigma")
sigma = lsl.Var.new_calc(jnp.exp, log_sigma, name="sigma")
mu = lsl.Var.new_calc(lambda X, beta: X * beta, X, beta, name="mu")
y = lsl.Var.new_obs(
    jnp.asarray(0.5 * x + rng.normal(scale=0.7, size=x.size)),
    lsl.Dist(tfd.Normal, mu, sigma),
    name="y",
)
model = lsl.Model(y)

Set the learning rate

Pass a configured Optax optimizer directly:

schedule = optax.exponential_decay(
    init_value=0.01, transition_steps=100, decay_rate=0.9
)
result = opt.LieselOptim(
    model,
    optimizers=optax.adam(schedule),
    loss_monitor="train_full_data",
    show_progress=False,
).fit()

result.position_min_monitor

{'beta': Array(0.43691176, dtype=float32, weak_type=True),
 'log_sigma': Array(-0.5869172, dtype=float32, weak_type=True)}

The schedule advances on optimizer updates, not epochs. With minibatches, it can advance several times per epoch. For a constant rate, use optax.adam(0.01). Minibatch fits can retain optimization noise. For precise point estimates, check convergence and consider a smaller rate or a final full-data fit.

Choose parameters

Give the slope and log noise scale different learning rates:

optimizers = [
    opt.Optimizer(["beta"], optax.adam(0.01)),
    opt.Optimizer(["log_sigma"], optax.adam(0.001)),
]

Pass this list as optimizers=optimizers to LieselOptim. Each optimizer updates its own parameters in list order on every batch. Their parameter names must not overlap. Parameters left out of the list stay fixed. Separate blocks each evaluate the loss and gradient, adding work.

L-BFGS must be the only optimizer because other parameter updates invalidate its cached objective and curvature history. Use optimizers="lbfgs" for all model parameters, or optimizers=[opt.LBFGS(["beta", "log_sigma"])] for a selected subset.

Choose the source

A prior may be attached to a weak parameter computed from a strong source variable. That prior remains in the default loss, including its derivatives through the weak parameter. Automatic parameter selection cannot decide which source to estimate and raises an informative error. Supply the strong names explicitly, for example optimizers=[opt.Optimizer(["source"], optax.adam(0.01))] or optimizers=[opt.LBFGS(["source"])]. The source does not need to be marked as a parameter. The weak parameter itself is recomputed during fitting.

Switch to L-BFGS

Use two separate fits to switch from Adam to L-BFGS. The optimizer option activate_after_epochs delays activation; it does not deactivate Adam. Keep the same parameter names in both fits:

first = opt.LieselOptim(
    model,
    optimizers=optax.adam(0.01),
    loss_monitor="train_full_data",
    show_progress=False,
).fit()
model.state = model.update_state(first.position_final, model.state)
result = opt.LieselOptim(
    model,
    optimizers="lbfgs",
    loss_monitor="train_full_data",
    show_progress=False,
).fit()

result.position_min_monitor

{'beta': Array(0.43742487, dtype=float32, weak_type=True),
 'log_sigma': Array(-0.5869661, dtype=float32, weak_type=True)}

Assign the returned model state before constructing the second fit: update_state leaves the model unchanged by default. L-BFGS starts with fresh optimizer state at Adam’s final parameter values and uses full-data batches here. Do not pass Adam’s checkpoint to the L-BFGS fit.

Save less history

Disable parameter history when you only need the final and best positions:

result = opt.LieselOptim(
    model,
    optimizers=optax.adam(0.01),
    loss_monitor="train_full_data",
    show_progress=False,
    save_position_history=False,
).fit()

History reserves memory for the maximum epoch budget before fitting, even when early stopping ends the run sooner. Its parameter storage costs approximately epochs * total_parameter_bytes: one million float32 parameters across 1,000 epochs take about 4 GB. Disabling it retains scalar losses and final/best positions.

Use build_engine() for settings beyond the wrapper’s arguments. See OptimEngine for all settings, or Pause and resume to pause and resume a fit.