Schedulers

NWAVE schedulers help you apply constraints progressively during training instead of enforcing them as hard constraints from the first epoch.

This is especially useful for hardware-oriented training, where constraints depend on the target chip architecture. Both H1v1 and H1v2 introduce chip-specific constraints, such as those described in the H1v1 constraints and H1v2 constraints, which can negatively impact convergence if enforced too aggressively from the start.

In both chips, sign-topology and weight-range constraints can hurt convergence if made too strict too early.

At the same time, there is a scheduler for surrogate parameters (that is not an hardware constraint but is in common with any spike model trained with gradient) that can be used in any layer of the SDK.

SignAnnealing

Sign Weight Magnitude Topology Scheduler

SignAnnealing progressively hardens row-block sign topology while keeping magnitude bounded to hardware range.

Import with:

from nwavesdk.optim import SignAnnealing

Why use it

Hardware imposes block-of-5 sign coherence (see H1v1/H1v2 chip constraints pages above).
Hardware also requires weight ranges ([-0.9, 0.9] for H1v1, [-1.66, 1.66] for H1v2).
Training with an immediately hard topology can make optimization unstable.
SignAnnealing starts with soft signs and progressively hardens them through alpha.

Internally it parametrizes 2D weight / recurrent_weights tensors with a row-block constrained parametrization (block size fixed to 5). Frontend layers are skipped.

How to use

from nwavesdk.optim import SignAnnealing

net = MySNN(...)
sign_annealer = SignAnnealing(
    net,
    total_epochs=epochs,
    alpha_start=0.5,
    alpha_end=20.0,
)

opt = torch.optim.Adam(net.parameters(), lr=lr)

for ep in range(1, epochs + 1):
    sign_annealer.step(net, ep)
    # forward, loss, backward, optimizer step

Warning

Create SignAnnealing before creating the optimizer. It registers weight parametrizations; if the optimizer is built first, it may miss the effective trainable parameters used by the scheduler.

Note

SignAnnealing is independent from your LR scheduler (for example ReduceLROnPlateau): you can use both together in the same loop.

Parameters

Constructor: SignAnnealing(net, total_epochs, alpha_start=0.5, alpha_end=10.0)

Parameter	Type	Default	Description
`net`	`torch.nn.Module`	required	Model where row-block parametrization is attached.
`total_epochs`	`int`	required	Number of epochs used by the alpha schedule.
`alpha_start`	`float`	`0.5`	Initial soft-sign sharpness. Lower = softer sign assignment.
`alpha_end`	`float`	`10.0`	Final sharpness. Higher = harder sign assignment near deployable topology.

Step call: step(net, epoch)

Parameter	Type	Default	Description
`net`	`torch.nn.Module`	required	Same model used at construction.
`epoch`	`int`	required	Current epoch index used to compute scheduled `alpha`.

SurrogateScheduler

SurrogateScheduler changes surrogate-gradient sharpness over training.

Import with:

from nwavesdk.optim import SurrogateScheduler

Why use it

Surrogate gradients control only backward dynamics (forward spiking stays hard-threshold).
Early training often benefits from smoother surrogate gradients.
Later training can benefit from sharper gradients to better align with hard spike behavior.
This scheduler lets you control that transition with reproducible policies.

Supported targets:

a full model (all modules exposing .spike_grad)
a single layer/module exposing .spike_grad
a surrogate object directly
a list/tuple of the above

How to use

from nwavesdk.optim import SurrogateScheduler

surrogate_sched = SurrogateScheduler(
    target=net,
    policy="cosine",
    start_value=10.0,
    end_value=30.0,
    total_steps=epochs,
    min_value=5.0,
    max_value=40.0,
)

for ep in range(epochs):
    current_surrogate_value = surrogate_sched.step(ep)
    # forward, loss, backward, optimizer step

Warning

All selected surrogates must expose the same scheduled parameter type (slope or alpha). Mixing types in one scheduler instance raises an error.

Warning

For policy="exp", both start_value and end_value must be strictly positive.

Note

Instantiate SurrogateScheduler before the training loop starts. Unlike SignAnnealing, it does not register new model parameters, so optimizer creation order is not critical.

Parameters

Constructor: SurrogateScheduler(target, policy="constant", start_value=None, end_value=None, total_steps=100, step_size=10, gamma=2.0, min_value=None, max_value=None, last_step=-1)

Parameter	Type	Default	Description
`target`	`nn.Module / surrogate / sequence`	required	What to schedule (model/layer/surrogate/list).
`policy`	`str`	`"constant"`	Schedule type: `constant`, `linear`, `exp`, `cosine`, `step`.
`start_value`	`float or None`	`None`	Initial value. If `None`, current surrogate value is used.
`end_value`	`float or None`	`None`	Final value for `linear`/`cosine`/`exp`. If `None`, equals `start_value`.
`total_steps`	`int`	`100`	Number of scheduler steps used to span the schedule.
`step_size`	`int`	`10`	Step interval for `policy="step"`.
`gamma`	`float`	`2.0`	Multiplicative factor for `step` policy (`start * gamma^k`).
`min_value`	`float or None`	`None`	Optional lower clip bound on scheduled value.
`max_value`	`float or None`	`None`	Optional upper clip bound on scheduled value.
`last_step`	`int`	`-1`	Internal step index for resume/checkpoint workflows.

Step call: step(step=None)

Parameter	Type	Default	Description
`step`	`int or None`	`None`	If provided, set explicit step index; otherwise increments internal counter.