omnibias.torch.training.k_scheduler
PyTorch activation derivatives, trainable operator layers, network jets, and direct differential-operator contractions.
import omnibias.torch.training.k_schedulerModule contract & mathematical conventions
Plateau-triggered K-growth scheduler.
Watches a scalar metric (typically validation loss); when it has not improved
by min_delta for patience consecutive epochs, the scheduler calls
grow(strategy) on every GrowableOperatorMultiBiasUnit in the
model. After growth, an LR-boost window is opened so freshly-activated bias
columns can warm up faster than the rest of the network.
Usage::
sched = KGrowthScheduler( model, patience=10, min_delta=1e-3, max_K=8, strategy="pair", lr_boost_factor=10.0, lr_boost_epochs=5, ) base_lr = optimizer.param_groups[0]["lr"]
for epoch in range(num_epochs): train_one_epoch(...) val_loss = evaluate(...) boost = sched.step(val_loss, epoch=epoch) for pg in optimizer.param_groups: pg["lr"] = base_lr * boost
Shared formulas do not imply universal floating-point identity. Verify dtype, shape, derivative-order support, and the stated numerical assumptions for the path you use.
AnchorProvider
omnibias.torch.training.AnchorProviderSource ↗AnchorProvider = Callable[[str, GrowableOperatorMultiBiasUnit], 'float | torch.Tensor | None']Type alias: an anchor provider returns the bias value(s) at which the new pair should be anchored. Receives the unit's name + the unit itself, so it can inspect existing biases to decide where to place the new pair.
Read the implementation
AnchorProvider = Callable[[str, GrowableOperatorMultiBiasUnit], "float | torch.Tensor | None"]Other public import paths
omnibias.torch.training.AnchorProvideromnibias.torch.training.k_scheduler.AnchorProvider
GrowthEvent
omnibias.torch.training.GrowthEventSource ↗class GrowthEventDataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.
One growth event recorded by the scheduler.
Parameters
| Name | Type & default | Contract |
|---|---|---|
epochdataclass field | intRequired | See the operation contract and implementation below. |
module_namedataclass field | strRequired | See the operation contract and implementation below. |
new_active_Kdataclass field | intRequired | See the operation contract and implementation below. |
strategydataclass field | strRequired | See the operation contract and implementation below. |
anchordataclass field | float | NoneDefault: None | See the operation contract and implementation below. |
Fields & methods
epochfield
epoch: intDeclared int field on GrowthEvent. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.
View source ↗module_namefield
module_name: strDeclared str field on GrowthEvent. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.
View source ↗new_active_Kfield
new_active_K: intDeclared int field on GrowthEvent. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.
View source ↗strategyfield
strategy: strDeclared str field on GrowthEvent. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.
View source ↗anchorfield
anchor: float | None = NoneDeclared float | None field on GrowthEvent. Default expression: None.
Read the implementation
class GrowthEvent:
"""One growth event recorded by the scheduler."""
epoch: int
module_name: str
new_active_K: int
strategy: str
anchor: float | None = NoneOther public import paths
omnibias.torch.training.GrowthEventomnibias.torch.training.k_scheduler.GrowthEvent
KGrowthScheduler
omnibias.torch.training.KGrowthSchedulerSource ↗class KGrowthSchedulerDataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.
Plateau-triggered K-growth controller.
Parameters
model : nn.Module Any module containing GrowableOperatorMultiBiasUnit children.
patience : int, default 10 Epochs of no-improvement before growth is triggered.
min_delta : float, default 1e-3 Minimum *relative* improvement counted as progress (loss must drop below best * (1 - min_delta) to reset patience).
max_K : int, default 8 Cap on each unit's active_K. The scheduler stops growing units that have reached this cap.
strategy : "pair" (default) or "saturate"
lr_boost_factor : float, default 10.0 Multiplicative LR boost applied for lr_boost_epochs after each growth event. The training loop is responsible for reading the scheduler's current_lr_multiplier and applying it to the optimizer's lr.
lr_boost_epochs : int, default 5
cooldown_epochs : int, default 5 Minimum epochs between successive growth events; gives the new biases time to settle before the next plateau trigger fires.
Parameters
| Name | Type & default | Contract |
|---|---|---|
modeldataclass field | nn.ModuleRequired | nn.Module Any module containing GrowableOperatorMultiBiasUnit children. |
patiencedataclass field | intDefault: 10 | int, default 10 Epochs of no-improvement before growth is triggered. |
min_deltadataclass field | floatDefault: 0.001 | float, default 1e-3 Minimum *relative* improvement counted as progress (loss must drop below best * (1 - min_delta) to reset patience). |
max_Kdataclass field | intDefault: 8 | int, default 8 Cap on each unit's active_K. The scheduler stops growing units that have reached this cap. |
strategydataclass field | GrowStrategyDefault: 'pair' | "pair" (default) or "saturate" |
lr_boost_factordataclass field | floatDefault: 10.0 | float, default 10.0 Multiplicative LR boost applied for lr_boost_epochs after each growth event. The training loop is responsible for reading the scheduler's current_lr_multiplier and applying it to the optimizer's lr. |
lr_boost_epochsdataclass field | intDefault: 5 | int, default 5 |
cooldown_epochsdataclass field | intDefault: 5 | int, default 5 Minimum epochs between successive growth events; gives the new biases time to settle before the next plateau trigger fires. |
anchor_providerdataclass field | AnchorProvider | NoneDefault: None | See the operation contract and implementation below. |
best_lossdataclass field | floatDefault: float('inf') | See the operation contract and implementation below. |
epochs_since_improvedataclass field | intDefault: 0 | See the operation contract and implementation below. |
cooldown_remainingdataclass field | intDefault: 0 | See the operation contract and implementation below. |
boost_remainingdataclass field | intDefault: 0 | See the operation contract and implementation below. |
growth_historydataclass field | list[GrowthEvent]Default: field(default_factory=list) | See the operation contract and implementation below. |
Executed regression · scheduler rejects models without growable units
Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.
def test_scheduler_rejects_models_without_growable_units() -> None:
plain = nn.Linear(2, 2)
with pytest.raises(ValueError):
KGrowthScheduler(plain)python -m pytest "packages/omnibias-torch/tests/test_growable.py::test_scheduler_rejects_models_without_growable_units" -qRegression example · scheduler respects max K cap
This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.
def test_scheduler_respects_max_K_cap() -> None:
m = _toy_growable_model(K_max=6)
sched = KGrowthScheduler(
m, patience=1, max_K=4, strategy="pair", lr_boost_epochs=0, cooldown_epochs=0
)
for ep in range(20):
sched.step(1.0, epoch=ep)
assert m.act.active_K <= 4, f"max_K=4 cap violated; active_K={m.act.active_K}"python -m pytest packages/omnibias-torch/tests/test_growable.py -k test_scheduler_respects_max_K_cap -qFields & methods
modelfield
model: nn.ModuleAny module containing GrowableOperatorMultiBiasUnit children.
min_deltafield
min_delta: float = 0.001Minimum *relative* improvement counted as progress (loss must drop below best * (1 - min_delta) to reset patience).
max_Kfield
max_K: int = 8Cap on each unit's active_K. The scheduler stops growing units that have reached this cap.
lr_boost_factorfield
lr_boost_factor: float = 10.0Multiplicative LR boost applied for lr_boost_epochs after each growth event. The training loop is responsible for reading the scheduler's current_lr_multiplier and applying it to the optimizer's lr.
cooldown_epochsfield
cooldown_epochs: int = 5Minimum epochs between successive growth events; gives the new biases time to settle before the next plateau trigger fires.
anchor_providerfield
anchor_provider: AnchorProvider | None = NoneOptional callable invoked at every growth event to decide where the new pair lands. Signature: (module_name, growable_unit) -> float | Tensor | None. Returning None falls back to anchoring at the unit's first existing bias (the default grow() behaviour). Recommended pattern: sample a value from the empirical distribution of pre-activations seen by growable_unit so the new pair lands inside the activation's informative range (avoiding the "dead-tail" failure mode where new biases drift to where sigma' vanishes).
epochs_since_improvefield
epochs_since_improve: int = 0Declared int field on KGrowthScheduler. Default expression: 0.
cooldown_remainingfield
cooldown_remaining: int = 0Declared int field on KGrowthScheduler. Default expression: 0.
boost_remainingfield
boost_remaining: int = 0Declared int field on KGrowthScheduler. Default expression: 0.
growth_historyfield
growth_history: list[GrowthEvent] = field(default_factory=list)Declared list[GrowthEvent] field on KGrowthScheduler. Default expression: field(default_factory=list).
__post_init__method
__post_init__(self) -> NoneSee the class contract above and the exact method implementation below.
Returns: None
Validation & failure cases
- When
not units, raisesValueError('KGrowthScheduler: the model contains no GrowableOperatorMultiBiasUnit children; nothing to grow.').
Executed regression · scheduler rejects models without growable units
Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.
def test_scheduler_rejects_models_without_growable_units() -> None:
plain = nn.Linear(2, 2)
with pytest.raises(ValueError):
KGrowthScheduler(plain)python -m pytest "packages/omnibias-torch/tests/test_growable.py::test_scheduler_rejects_models_without_growable_units" -qMethod implementation
def __post_init__(self) -> None:
# Validate that the model has at least one growable unit.
units = self._growable_units()
if not units:
raise ValueError(
"KGrowthScheduler: the model contains no GrowableOperatorMultiBiasUnit "
"children; nothing to grow."
)current_lr_multiplierproperty
current_lr_multiplier(self) -> floatThe multiplier the training loop should apply to the base LR.
Returns: float
Executed regression · scheduler respects max K cap
Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.
def test_scheduler_respects_max_K_cap() -> None:
m = _toy_growable_model(K_max=6)
sched = KGrowthScheduler(
m, patience=1, max_K=4, strategy="pair", lr_boost_epochs=0, cooldown_epochs=0
)
for ep in range(20):
sched.step(1.0, epoch=ep)
assert m.act.active_K <= 4, f"max_K=4 cap violated; active_K={m.act.active_K}"python -m pytest "packages/omnibias-torch/tests/test_growable.py::test_scheduler_respects_max_K_cap" -qMethod implementation
def current_lr_multiplier(self) -> float:
"""The multiplier the training loop should apply to the base LR."""
return self.lr_boost_factor if self.boost_remaining > 0 else 1.0stepmethod
step(self, metric: float, *, epoch: int) -> floatUpdate internal state with the latest metric value (lower-is-better).
Returns the LR multiplier the training loop should apply for the
*next* epoch. The multiplier equals lr_boost_factor for
lr_boost_epochs epochs after a growth event, and 1.0
otherwise.
Patience semantics: the very first call establishes the baseline
(counts neither as improvement nor as plateau). From the second
call onwards, epochs_since_improve is incremented unless the
metric drops below best_loss * (1 - min_delta).
Parameters
| Name | Type & default | Contract |
|---|---|---|
metricpositional-or-keyword | floatRequired | See the operation contract and implementation below. |
epochkeyword-only | intRequired | See the operation contract and implementation below. |
Returns: float
Executed regression · scheduler respects max K cap
Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.
def test_scheduler_respects_max_K_cap() -> None:
m = _toy_growable_model(K_max=6)
sched = KGrowthScheduler(
m, patience=1, max_K=4, strategy="pair", lr_boost_epochs=0, cooldown_epochs=0
)
for ep in range(20):
sched.step(1.0, epoch=ep)
assert m.act.active_K <= 4, f"max_K=4 cap violated; active_K={m.act.active_K}"python -m pytest "packages/omnibias-torch/tests/test_growable.py::test_scheduler_respects_max_K_cap" -qMethod implementation
def step(self, metric: float, *, epoch: int) -> float:
"""Update internal state with the latest metric value (lower-is-better).
Returns the LR multiplier the training loop should apply for the
*next* epoch. The multiplier equals ``lr_boost_factor`` for
``lr_boost_epochs`` epochs after a growth event, and ``1.0``
otherwise.
Patience semantics: the very first call establishes the baseline
(counts neither as improvement nor as plateau). From the second
call onwards, ``epochs_since_improve`` is incremented unless the
metric drops below ``best_loss * (1 - min_delta)``.
"""
# Decrement boost first; the multiplier we return at the end of
# this call applies to the *next* training step.
if self.boost_remaining > 0:
self.boost_remaining -= 1
if not self._initialized:
self.best_loss = metric
self.epochs_since_improve = 0
self._initialized = True
else:
improved = metric < self.best_loss * (1.0 - self.min_delta)
if improved:
self.best_loss = metric
self.epochs_since_improve = 0
else:
self.epochs_since_improve += 1
# Cooldown blocks growth while still positive; decrement happens
# *after* the readiness check so a fresh grow with cooldown=N
# blocks exactly the next N step() calls.
ready = (
self.epochs_since_improve >= self.patience
and self.cooldown_remaining == 0
and self._has_pair_capacity_or_saturate_capacity()
)
grew_this_step = False
if ready:
grew_this_step = self._grow_all_eligible(epoch=epoch)
if grew_this_step:
self.epochs_since_improve = 0
self.boost_remaining = self.lr_boost_epochs
self.cooldown_remaining = self.cooldown_epochs
if not grew_this_step and self.cooldown_remaining > 0:
self.cooldown_remaining -= 1
return self.current_lr_multiplierreset_bestmethod
reset_best(self) -> NoneForget the best-so-far loss; useful after dataset switch.
Returns: None
Method implementation
def reset_best(self) -> None:
"""Forget the best-so-far loss; useful after dataset switch."""
self.best_loss = float("inf")
self.epochs_since_improve = 0
self._initialized = Falseactive_K_summarymethod
active_K_summary(self) -> dict[str, int]Return {module_name: active_K} for all growable units.
Returns: dict[str, int]
Method implementation
def active_K_summary(self) -> dict[str, int]:
"""Return ``{module_name: active_K}`` for all growable units."""
return {name: mod.active_K for name, mod in self._growable_units()}total_growth_eventsmethod
total_growth_events(self) -> intSee the class contract above and the exact method implementation below.
Returns: int
Executed regression · scheduler resets patience on improvement
Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.
def test_scheduler_resets_patience_on_improvement() -> None:
m = _toy_growable_model(K_max=6)
sched = KGrowthScheduler(m, patience=3, max_K=6, strategy="pair")
sched.step(1.0, epoch=0) # baseline
sched.step(1.0, epoch=1) # plateau 1
sched.step(1.0, epoch=2) # plateau 2
sched.step(0.5, epoch=3) # improvement -> reset
sched.step(0.5, epoch=4)
sched.step(0.5, epoch=5)
sched.step(0.5, epoch=6) # 3 consecutive plateau steps at 0.5 -> grow
assert m.act.active_K == 3, "should grow only after the second plateau"
assert sched.total_growth_events() == 1python -m pytest "packages/omnibias-torch/tests/test_growable.py::test_scheduler_resets_patience_on_improvement" -qMethod implementation
def total_growth_events(self) -> int:
return len(self.growth_history)Read the implementation
class KGrowthScheduler:
"""Plateau-triggered K-growth controller.
Parameters
----------
model : nn.Module
Any module containing :class:`GrowableOperatorMultiBiasUnit` children.
patience : int, default 10
Epochs of no-improvement before growth is triggered.
min_delta : float, default 1e-3
Minimum *relative* improvement counted as progress
(loss must drop below ``best * (1 - min_delta)`` to reset patience).
max_K : int, default 8
Cap on each unit's ``active_K``. The scheduler stops growing units
that have reached this cap.
strategy : ``"pair"`` (default) or ``"saturate"``
lr_boost_factor : float, default 10.0
Multiplicative LR boost applied for ``lr_boost_epochs`` after each
growth event. The training loop is responsible for reading the
scheduler's ``current_lr_multiplier`` and applying it to the
optimizer's ``lr``.
lr_boost_epochs : int, default 5
cooldown_epochs : int, default 5
Minimum epochs between successive growth events; gives the new
biases time to settle before the next plateau trigger fires.
"""
model: nn.Module
patience: int = 10
min_delta: float = 1e-3
max_K: int = 8
strategy: GrowStrategy = "pair"
lr_boost_factor: float = 10.0
lr_boost_epochs: int = 5
cooldown_epochs: int = 5
anchor_provider: AnchorProvider | None = None
"""Optional callable invoked at every growth event to decide where the
new pair lands. Signature: ``(module_name, growable_unit) -> float |
Tensor | None``. Returning ``None`` falls back to anchoring at the
unit's first existing bias (the default ``grow()`` behaviour).
Recommended pattern: sample a value from the empirical distribution of
pre-activations seen by ``growable_unit`` so the new pair lands inside
the activation's informative range (avoiding the "dead-tail" failure
mode where new biases drift to where ``sigma'`` vanishes)."""
# Internal state.
best_loss: float = float("inf")
epochs_since_improve: int = 0
cooldown_remaining: int = 0
boost_remaining: int = 0
growth_history: list[GrowthEvent] = field(default_factory=list)
_initialized: bool = False
def __post_init__(self) -> None:
# Validate that the model has at least one growable unit.
units = self._growable_units()
if not units:
raise ValueError(
"KGrowthScheduler: the model contains no GrowableOperatorMultiBiasUnit "
"children; nothing to grow."
)
# ----- public API ------------------------------------------------------
@property
def current_lr_multiplier(self) -> float:
"""The multiplier the training loop should apply to the base LR."""
return self.lr_boost_factor if self.boost_remaining > 0 else 1.0
def step(self, metric: float, *, epoch: int) -> float:
"""Update internal state with the latest metric value (lower-is-better).
Returns the LR multiplier the training loop should apply for the
*next* epoch. The multiplier equals ``lr_boost_factor`` for
``lr_boost_epochs`` epochs after a growth event, and ``1.0``
otherwise.
Patience semantics: the very first call establishes the baseline
(counts neither as improvement nor as plateau). From the second
call onwards, ``epochs_since_improve`` is incremented unless the
metric drops below ``best_loss * (1 - min_delta)``.
"""
# Decrement boost first; the multiplier we return at the end of
# this call applies to the *next* training step.
if self.boost_remaining > 0:
self.boost_remaining -= 1
if not self._initialized:
self.best_loss = metric
self.epochs_since_improve = 0
self._initialized = True
else:
improved = metric < self.best_loss * (1.0 - self.min_delta)
if improved:
self.best_loss = metric
self.epochs_since_improve = 0
else:
self.epochs_since_improve += 1
# Cooldown blocks growth while still positive; decrement happens
# *after* the readiness check so a fresh grow with cooldown=N
# blocks exactly the next N step() calls.
ready = (
self.epochs_since_improve >= self.patience
and self.cooldown_remaining == 0
and self._has_pair_capacity_or_saturate_capacity()
)
grew_this_step = False
if ready:
grew_this_step = self._grow_all_eligible(epoch=epoch)
if grew_this_step:
self.epochs_since_improve = 0
self.boost_remaining = self.lr_boost_epochs
self.cooldown_remaining = self.cooldown_epochs
if not grew_this_step and self.cooldown_remaining > 0:
self.cooldown_remaining -= 1
return self.current_lr_multiplier
def reset_best(self) -> None:
"""Forget the best-so-far loss; useful after dataset switch."""
self.best_loss = float("inf")
self.epochs_since_improve = 0
self._initialized = False
# ----- introspection ---------------------------------------------------
def active_K_summary(self) -> dict[str, int]:
"""Return ``{module_name: active_K}`` for all growable units."""
return {name: mod.active_K for name, mod in self._growable_units()}
def total_growth_events(self) -> int:
return len(self.growth_history)
# ----- internals -------------------------------------------------------
def _growable_units(self) -> list[tuple[str, GrowableOperatorMultiBiasUnit]]:
return [
(name, mod)
for name, mod in self.model.named_modules()
if isinstance(mod, GrowableOperatorMultiBiasUnit)
]
def _strategy_increment(self) -> int:
return 2 if self.strategy == "pair" else 1
def _has_pair_capacity_or_saturate_capacity(self) -> bool:
"""At least one unit can still grow under this scheduler's max_K
(and per-unit K_max) without overshooting."""
inc = self._strategy_increment()
for _, mod in self._growable_units():
if mod.active_K + inc <= min(self.max_K, mod.K_max):
return True
return False
def _grow_all_eligible(self, *, epoch: int) -> bool:
inc = self._strategy_increment()
grew_any = False
for name, mod in self._growable_units():
if mod.active_K + inc > min(self.max_K, mod.K_max):
continue
anchor: float | torch.Tensor | None = None
if self.anchor_provider is not None and self.strategy == "pair":
try:
anchor = self.anchor_provider(name, mod)
except Exception as exc: # noqa: BLE001
# An anchor provider failure should not abort training;
# log and fall back to the default anchor (anchor=None).
_LOGGER.warning(
"anchor_provider failed for %r: %r; falling back to "
"default anchor.",
name,
exc,
)
anchor = None
try:
if anchor is not None:
mod.grow(strategy=self.strategy, anchor_value=anchor)
else:
mod.grow(strategy=self.strategy)
except (RuntimeError, ValueError):
# Hit per-unit K_max during grow, or strategy not supported
# for this activation. Skip the unit; do not abort.
continue
anchor_scalar: float | None = None
if isinstance(anchor, torch.Tensor):
anchor_scalar = float(anchor.float().mean().item())
elif anchor is not None:
anchor_scalar = float(anchor)
self.growth_history.append(
GrowthEvent(
epoch=epoch,
module_name=name,
new_active_K=mod.active_K,
strategy=self.strategy,
anchor=anchor_scalar,
)
)
grew_any = True
return grew_anyOther public import paths
omnibias.torch.training.KGrowthScheduleromnibias.torch.training.k_scheduler.KGrowthScheduler