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 use adam). Stops on an exponentially-smoothed full-batch loss with an absolute-iteration window (plateau_lag * check_every iterations) – decoupled from batch size (the epoch-based window made small batches fragile; findings in a separate research repo, maintainer-local). use_jit jits 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 default flat_draw then slices the on-device sample / data (a custom draw should 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]