mc_sgd#
- t3toolbox.backend.optimizers.mc_sgd(problem, x0, rng, batch, draw=None, max_iter=3000, check_every=25, smooth_tau=2.0, plateau_lag=4, use_jit=False)#
def mc_sgd( problem: Problem, # the fixed-rank least-squares problem x0: Tangent, # initial cores (U, G) rng, # np.random.Generator -- passed to the draw each step batch: int, # measurements per minibatch (for the default flat draw; ignored if draw given) draw: typ.Optional[typ.Callable] = None, # custom draw(rng)->(sample_B,data_B); None = flat default max_iter: int = 3000, check_every: int = 25, # iterations between full-batch loss checks (the absolute-iteration window) smooth_tau: float = 2.0, # EMA timescale of the loss check, in checks plateau_lag: int = 4, # stop when the smoothed loss rises over this many checks use_jit: bool = False, # jit the per-step kernel (gradient + Cauchy step + retract) when jax ) -> typ.Tuple[Tangent, dict]: # (x_cores, stats)
Manifold Cauchy SGD (T4S §5.3.2): minibatch + the tuning-free Cauchy step
α = ‖g‖² / ‖𝒥g‖²– no learning rate. Intended for the manifold geometry (its bounded retraction makes the raw Cauchy step stable; on the additive corewise chart useadam). Stops on an exponentially-smoothed full-batch loss with an absolute-iteration window (plateau_lag * check_everyiterations) – decoupled from batch size (the epoch-based window made small batches fragile; findings in a separate research repo, maintainer-local).use_jitjits the per-step kernel (the host loop draws minibatches; the full-batch stop check stays on the host), auto-converting numpy inputs to jax first (_prepare_jit_inputs()) – so the defaultflat_drawthen slices the on-devicesample/data(a customdrawshould return jax minibatches to stay on device).- Parameters:
problem (Problem)
x0 (Tangent)
batch (int)
draw (t3toolbox.backend.common.typ.Optional[t3toolbox.backend.common.typ.Callable])
max_iter (int)
check_every (int)
smooth_tau (float)
plateau_lag (int)
use_jit (bool)
- Return type:
t3toolbox.backend.common.typ.Tuple[Tangent, dict]