Schedules: what varies with the step

A schedule varies one quantity per optimizer step: the learning rate, a decay coefficient, another rule field, a Device field read inside a traced hook, or a gradient rule field such as SAM's rho. Everything else is fixed at setup.

The contract

schedules(e) returns a NamedTuple; the default is (;), meaning everything constant. Every entry is a factory of the horizon: the framework calls f(total) once, at setup, with the total number of optimizer steps, then calls what it returned once per optimizer step. A bare Number normalizes to _ -> (_ -> value), a constant.

schedules(e::MyExp) = (;    eta        = total -> OneCycle(total, 1f-3),               # factory of the horizon    aux_weight = _ -> (t -> max(0f0, 1f0 - t / 5000)),         # ignores the horizon    lambda     = 1f-4,                                         # a bare Number is a constant)

eta and lambda are rule fields and aux_weight is a Device field on MyExp; the next sections say how each key resolves. OneCycle(total, 1f-3) is the schedule called once per step, aux_weight's factory ignores the horizon, and lambda is constant for the whole run.

Step and horizon

step is the optimizer step, not the micro-batch. With accum = N they differ by a factor of N. The horizon is total = max_epochs * div(length(train), accum), where length(train) is the number of training batches per epoch. The division is exact because a remainder is a setup error: the training loader must drop its partial final batch, and its batch count must divide by accum.

Schedules run on the host, so ordinary Julia control flow is fine inside them: an if, a for, a call into any library. A schedule need not be a pure function of the step, but one that closes over training state does not resume exactly. To react to validation, use Plateau, which does.

The two namespaces

A schedule key is either a Device field name or a rule field name, resolved against the union of the two. An opt key is applied to the optimizer; a device key is written into the experiment and reaches traced code as e.field.

device and opt are reserved at the top level and name the two namespaces; every other top-level key is unqualified and must resolve to exactly one of them. A key matching neither is an error naming the fix. A key matching more than one is refused, and the error shows the qualified forms and the bind form, so you never have to rename the field. Collisions are common: lambda is a natural name for a decay coefficient and for a loss weight, and eta, rho, and epsilon are plausible model hyperparameters.

schedules = (; eta = total -> OneCycle(total, 1f-3),               device = (; lambda = _ -> t -> 0.5f0),    # e.lambda               opt     = (; lambda = _ -> t -> 1f-4))     # the decay coefficient

Rule field names are the rule's actual field names, so the learning-rate key is eta, not lr.

A third reserved namespace, grad, holds the fields of the grad_rule, such as SAM's and ASAM's rho (ASAM's eta is not schedulable, so an unqualified eta stays the learning rate). An unqualified rho resolves to it when nothing else has that name; with Optimisers.Momentum, whose field is also rho, it is refused and you qualify it. The value logs as schedule/grad/rho.

schedules = (; grad = (; rho = total -> t -> 0.05f0 * min(1, t / 500)))

Binding by hand: bind

Every namespace above binds by name, which is where collisions come from. bind holds names of your choosing that resolve against nothing; you place the values yourself, in a Level 2 optimizer(e, group, hp) and in grad_rule(e, hp), where they arrive as hp.bind.<name>:

schedules(::MyExp) = (;    eta  = total -> OneCycle(total, 1f-3),    bind = (; mom_ramp = _ -> t -> min(1f0, t / 1000), sam_warm = _ -> t -> min(1f0, t / 500)),)ReactantNitro.optimizer(::MyExp, group::Symbol, hp) =    Optimisers.Momentum(hp.eta, 0.85f0 + 0.1f0 * hp.bind.mom_ramp)ReactantNitro.grad_rule(::MyExp, hp) = SAM(0.05f0 * hp.bind.sam_warm)

A value computed from hp, as here, follows its schedule and never recompiles: setup probes each hook with two values of each bind key, records which rule fields change, and the framework moves those fields to the device every step. A bind key that changes nothing is a setup error. The binding report lists where each key landed, and the train line logs schedule/bind/<name>. bind covers the optimizer and the gradient rule; a Device field is the experiment's own named slot and needs no binding.

Scheduling lambda has one extra rule: opt.lambda together with per-group lambda accessors is a setup error naming both, because the schedule supplies the decay coefficient for every group. A path-bound opt.<group>.lambda (next section) conflicts only with that group's accessor. A scheduled lambda keeps the Decay tail in the chain even while its value is zero, so a schedule that starts or ends at zero does not change the chain mid-run. eta has no such conflict: per-group learning_rate accessors are ratios that multiply an eta schedule (Optimization).

Per-group curves: path-bound opt keys

A bare opt key applies to every parameter group. A path-bound key, opt.<group>.<field>, applies to the named group only, so groups can follow different curves.

schedules(e::MyExp) = (;    opt = (;        eta      = total -> OneCycle(total, 1f-3),               # bare: groups without a path key        backbone = (; eta = total -> t -> 1f-3 * (1 - t / total)),  # path: the :backbone group    ))

For each group and field, the value comes from the first of these that exists:

1. the path key for that group, opt.<group>.<field>;
2. the bare key, opt.<field> or an unqualified <field> that resolves to the optimizer;
3. the group's accessor: learning_rate(e, Val(g)) for eta, lambda for lambda.

So above, :backbone decays linearly and every other group follows OneCycle. Any other field with no key keeps the rule's own value.

A path value goes through the same per-group scaling as a bare one:

# η_g(t) = value(t) * (learning_rate(e, Val(g)) / learning_rate(e))   # eta# λ_g(t) = η_g(t) * value(t)                                          # lambda# any other field: value(t), as given

so opt.backbone.eta = 1f-2 against a backbone ratio of 0.1 trains the backbone at 1f-3. To give a group an absolute curve, leave its learning_rate accessor at the default ratio of 1.

Path keys are qualified only: they live under opt, and an unqualified dotted key is a resolution error. Paths are one level (group.field), since parameter groups are flat. A path segment that is not a parameter group, and a field that is not schedulable in that group's chain, are setup errors naming the offender.

Manual mode uses the same syntax with different values: there are no groups or ratios, a path is a path into opt_state, and every value is absolute (Manual training).

Reacting to validation: Plateau

Plateau wraps any schedule and scales it down when a validation metric stops improving, as PyTorch's ReduceLROnPlateau does:

schedules(e::MyExp) = (;    eta = Plateau(total -> OneCycle(total, 1f-3); patience = 3),    device = (; aux_weight = Plateau(0.5f0; metric = :auc, mode = :max, factor = 0.5)),)

The value at step t is inner(t) * scale. The scale starts at 1. After each validation, more than patience epochs without improvement multiply it by factor, floored at min_scale, and then cooldown epochs pass before bad epochs count again. The rule and the defaults are PyTorch's (factor = 0.1, patience = 10, threshold = 1f-4 relative), so patience = 2 reduces at the third bad epoch. A bare number as inner is the plain ReduceLROnPlateau.

Every value is a pure function of the step and the run's validation history. The counters are stored in each checkpoint, so a resume continues the same sequence; a weights load starts fresh. Each key has its own counters. A reduction changes a value and never a type, so it recompiles nothing.

The validation line logs schedule/plateau/<key>, each reduction is reported with @info, and the binding report marks the key with its metric, mode, factor and patience. Setup requires a val split, and an experiment without metrics can only watch :loss.

A plateau watches the run's monitor, as the checkpointer and EarlyStopping do, unless given its own metric and mode, as aux_weight's is above. With a custom metrics, a metric name cannot be checked until validation first runs, so a misspelled one raises at the end of the first epoch, not at setup.

How schedules interact with Device values

Anything you want to vary inside loss, forward, or metrics must be a Device field. A scheduled device key is written into the experiment on each step, and traced code reads e.field. A GraphConst cannot be scheduled: its value is baked into the compiled graph and is part of the compile cache key, so a different value is a different program. A Host field is in neither namespace, so a schedule key naming one is a resolution error naming the fix.

A scheduled value is converted to the field's own type, so a schedule written t -> 0.25 * min(1, t / 2000), which returns Float64, works for a Device{Float32} field without the f0. After the conversion the framework asserts that the experiment's type did not change; a change would alter the compile key and recompile the gradient program on every step.

A constant and a schedule use the same device slot. A bare Number normalizes to a constant factory, so switching a field between a constant and a schedule costs no recompile.

The fields a rule may not schedule are declared by nonschedulable(::Type{R}), which defaults to ():

nonschedulable(::Type{<:Optimisers.Adam})     = (:beta, :epsilon)nonschedulable(::Type{<:Optimisers.RAdam})    = (:beta, :epsilon)nonschedulable(::Type{<:Optimisers.AdamW})    = (:beta, :epsilon, :couple)nonschedulable(::Type{<:Optimisers.ClipNorm}) = (:p, :throw)nonschedulable(::Type{<:Decay})               = (:anchor, :no_decay_mask)

The same declaration also keeps those fields off the device, so what is schedulable and what is device-resident cannot drift apart. Any parameter-sized rule field is non-schedulable, since scheduling it would upload a parameter-sized buffer every step; Decay's anchor and no_decay_mask are the shipped case. The scalar exclusions have their own reasons: beta is a Tuple, which a scalar schedule cannot produce; couple is a Bool; p is structural, and promoting it breaks the clip norm's dispatch; epsilon is a numerical floor rather than a tuning knob.

Declare one method per concrete rule type, never a Union: a more specific method shadows a Union method with no ambiguity warning, so a scalar schedule could be written into a Tuple field.

The binding report

The binding report lists every schedule key, what it resolved to, and whether it came from a keyword or the accessor, plus one row per parameter group with that group's rate, ratio, anchor, decay, rule, and parameter count. At Level 2 it also shows which rule field each scheduled value landed in.

The report is part of show(nitro): printing a handle shows what the run holds and where each configured value came from, in one table. Setup prints nothing itself, since the REPL already displays the handle the constructor returns. From a script, call display(nitro). The same text is passed to log_other!, so the run's record carries it.

Read it on the first run of any new config: it shows a group that came out empty, a ratio you did not intend, or a schedule that bound to the optimizer when you meant your experiment. It reports and never fails.

The schedule belongs to the experiment by default, and the schedules keyword on Nitro and train! overrides it for one run. The keyword replaces the accessor's table; it does not merge with it. To merge, say so:

train!(e; schedules = merge(schedules(e), (; eta = total -> OneCycle(total, 3f-4))))

Logging

The framework logs every varying scalar schedule on the train line, which is written every log_every_n_steps optimizer steps (default 1). Every key starts with schedule/, so none can collide with a train_metrics name.

• A Device field logs as schedule/<field>.
• A rule value logs as schedule/<group>/<field>, for each group whose chain uses it, at the value applied: eta and lambda after the per-group scaling above, so schedule/backbone/eta is the rate the backbone trained at.
• A group's lambda is also logged when only its eta varies, if the group decays, since its decay is η_g(t) times the coefficient.
• Manual mode logs each opt key's value as given, with the path joined by / (schedule/gen/eta).

Constants, non-finite values, and values that are not Real (an array, a device scalar) are left out; the binding report records the constants. To turn logging off, pass Nitro(e; autolog_schedules = false) or define autolog_schedules(e) = false.

Use the ecosystem

ParameterSchedulers.jl is the recommendation. ReactantNitro does not depend on it, not even as a weakdep: a schedule is any callable of the step, so there is no interface to satisfy.

The framework calls a schedule with the upcoming optimizer step, counting from 1. ParameterSchedulers schedules are 1-indexed, so s(t) lines up with no offset. Check this for any library you bring: a 0-indexed schedule needs t -> s(t - 1).

Do not check a library's index convention at step 1. Most schedules start on a flat part, so the first two steps can differ by less than Float32 eps and round to the same number. Check mid-curve, where the slope is steepest.

PyTorch's torch.optim.lr_scheduler.OneCycleLR curve can be reproduced exactly, but only by mapping every argument rather than accepting defaults:

At max_lr = 1e-3, 20 steps, pct_start = 0.3, div_factor = 25, final_div_factor = 1e4, the two curves agree to a maximum absolute difference of 2.2e-19, with the same peak step and the same endpoint. With OneCycle's own defaults the final value differs by four orders of magnitude, because the two libraries define the end of the cycle differently.

OneCycle cannot express a zero-length warmup. It asserts 0 < percent_start < 1, so a schedule that starts at max_lr and only anneals (pct_start = 0 in PyTorch) fails to construct. percent_start = 1e-6 is accepted but is not the same curve.

Every name on this page is documented on the API page. For the field markers behind device schedules, see Experiments; for eta, parameter groups, and decay, see Optimization.