Optimizer

The optimizer is the component that updates the model’s weights from their gradients during training. It defines the optimization algorithm and its hyperparameters.

Usage

The optimizer occupies the components.optimizer slot.

components:
  optimizer:
    name: adam
    kwargs:
      lr: 0.0002
      weight_decay: 0.00005

How it Works

When the run is assembled, the optimizer takes the model component’s parameters and wraps the corresponding optimization algorithm. It exposes the two operations the training loop drives: a step that applies an update, and a gradient reset.

Variants

Adam: adaptive per-parameter step sizes with decoupled weight decay.
SGD: stochastic gradient descent, optionally with momentum and Nesterov acceleration.

Registering a new optimizer

An optimizer is a subclass of Optimizer that declares a name and version, validates its configuration, performs one-time setup in build, binds to the model parameters on attach, and implements the two operations the training loop calls:

build(self) -> None: One-time setup, before the model parameters are available.
step(self) -> None: Apply one optimization update.
zero_grad(self, set_to_none=True) -> None: Clear the accumulated gradients.

The underlying algorithm is typically constructed in on_attach, where the model component (and therefore its parameters) is available.

from typing import Any, ClassVar

from icegraph.common.engine import ComponentKind
from icegraph.engine.components.optimizer import Optimizer, OptimizerFactory

from .config import MyOptimizerConfig

class MyOptimizer(Optimizer[MyOptimizerConfig]):
    name: ClassVar[str] = "my-optimizer"
    version: ClassVar[int] = 1

    @classmethod
    def validate_config(cls, config: dict[str, Any]) -> MyOptimizerConfig:
        return MyOptimizerConfig(**config)

    def build(self) -> None:
        ...  # one-time setup before the model is available

    def on_attach(self) -> None:
        model = self._ctx.components.require(ComponentKind.MODEL, required_by=type(self))
        ...  # build the optimizer from model.parameters()

    def step(self) -> None:
        ...

    def zero_grad(self, set_to_none: bool = True) -> None:
        ...

OptimizerFactory.register(MyOptimizer)