optim

type Adam ¶

type Adam[B tensor.Backend] struct {
	// contains filtered or unexported fields
}

Adam implements the Adam (Adaptive Moment Estimation) optimizer with optional decoupled weight decay (AdamW).

Adam combines ideas from RMSprop and momentum:

• Maintains exponential moving averages of gradients (first moment)
• Maintains exponential moving averages of squared gradients (second moment)
• Applies bias correction to compensate for initialization at zero

Update rule (AdamW when WeightDecay > 0):

m_t   = beta1 * m_{t-1} + (1-beta1) * gradient       // First moment
v_t   = beta2 * v_{t-1} + (1-beta2) * gradient²      // Second moment
m_hat = m_t / (1 - beta1^t)                           // Bias correction
v_hat = v_t / (1 - beta2^t)                           // Bias correction
param = param * (1 - lr * weightDecay)                // Decoupled weight decay
      - lr * m_hat / (sqrt(v_hat) + eps)              // Adaptive update

All parameter updates use pure tensor operations — no AsFloat32() / CPU readback in the hot path. Bias-correction scalars are O(1) CPU values per step (they depend only on the scalar timestep, not on tensor data).

If the backend implements CacheInvalidator, ClearInputBufferCache() is called at the end of Step() to invalidate stale weight-buffer cache entries.

Reference: "Adam: A Method for Stochastic Optimization" (Kingma & Ba, 2014)

Example:

optimizer := optim.NewAdam(model.Parameters(), optim.AdamConfig{
    LR:    0.001,
    Betas: [2]float32{0.9, 0.999},
    Eps:   1e-8,
})

for epoch := range epochs {
    loss := train_step(model, batch)
    grads := autodiff.Backward(loss, backend)
    optimizer.Step(grads)
    optimizer.ZeroGrad()
}

func NewAdam[B tensor.Backend](params []*nn.Parameter[B], config AdamConfig, backend B) *Adam[B]

func (a *Adam[B]) DiagMoments() [][2]*tensor.RawTensor

func (a *Adam[B]) GetLR() float32

func (a *Adam[B]) GetTimestep() int

func (a *Adam[B]) LoadStateDict(stateDict map[string]*tensor.RawTensor) error

func (a *Adam[B]) SetLR(lr float32)

func (a *Adam[B]) StateDict() map[string]*tensor.RawTensor

func (a *Adam[B]) Step(grads map[*tensor.RawTensor]*tensor.RawTensor)

func (a *Adam[B]) ZeroGrad()

type AdamConfig ¶

type AdamConfig struct {
	LR          float32    // Learning rate (default: 0.001)
	Betas       [2]float32 // Coefficients for computing running averages (default: [0.9, 0.999])
	Eps         float32    // Term for numerical stability (default: 1e-8)
	WeightDecay float32    // Decoupled weight-decay coefficient — AdamW style (default: 0.0)
}

AdamConfig holds configuration for Adam optimizer.

type CacheInvalidator ¶ added in v0.9.1

type CacheInvalidator interface {
	ClearInputBufferCache()
}

CacheInvalidator is an optional interface that backends implement when they maintain an input-buffer cache keyed by *RawTensor identity.

After an optimizer step the weight RawTensors stored in Parameters are replaced by new ones (via SetTensor). The old *RawTensor pointers no longer represent the current weights; any backend cache entry under those keys is stale. Calling ClearInputBufferCache() forces the next forward pass to re-upload the freshly computed weights from the new RawTensors.

Optimizers call this via a type-assertion at the end of Step() — backends that do not maintain such a cache simply don't implement the interface.

type Config ¶

type Config struct {
	LR float32 // Learning rate
}

Config is the base configuration for all optimizers.

type Optimizer ¶

type Optimizer interface {
	// Step applies gradient updates to all parameters.
	//
	// Takes a gradient map from Backward() and updates parameters in-place.
	// The gradient map should contain RawTensor -> gradient mapping.
	//
	// Example:
	//   grads := autodiff.Backward(loss, backend)
	//   optimizer.Step(grads)
	Step(grads map[*tensor.RawTensor]*tensor.RawTensor)

	// ZeroGrad clears all parameter gradients.
	//
	// This should be called before each backward pass to prevent
	// gradient accumulation from previous iterations.
	//
	// Example:
	//   optimizer.ZeroGrad()
	//   loss := model.Forward(...)
	//   grads := autodiff.Backward(loss, backend)
	ZeroGrad()

	// GetLR returns the current learning rate.
	//
	// Useful for monitoring and learning rate scheduling.
	GetLR() float32

	// StateDict returns the optimizer state for serialization.
	//
	// This includes optimizer-specific state like momentum buffers (SGD)
	// or moment estimates (Adam). Used for checkpoint saving.
	//
	// Returns a map from state name to RawTensor.
	// State names follow the pattern: "{state_type}.{param_index}"
	// For example: "velocity.0", "m.0", "v.0"
	StateDict() map[string]*tensor.RawTensor

	// LoadStateDict loads optimizer state from serialization.
	//
	// Restores optimizer-specific state from a checkpoint. The state
	// dictionary should match the structure returned by StateDict().
	//
	// Parameters:
	//   - stateDict: Map from state name to RawTensor
	//
	// Returns an error if the state dictionary is invalid.
	LoadStateDict(stateDict map[string]*tensor.RawTensor) error
}

Optimizer is the base interface for all optimization algorithms.

Optimizers update model parameters based on computed gradients to minimize the loss function during training.

All optimizers must implement:

• Step: Apply gradient updates to parameters
• ZeroGrad: Clear gradients before next iteration
• GetLR: Get current learning rate (for monitoring/scheduling)
• StateDict: Export optimizer state for checkpoints
• LoadStateDict: Import optimizer state from checkpoints

type SGD ¶

type SGD[B tensor.Backend] struct {
	// contains filtered or unexported fields
}

SGD implements Stochastic Gradient Descent optimizer with optional momentum and L2 weight decay.

Update rule without momentum:

param = param * (1 - lr * weightDecay) - lr * gradient

Update rule with momentum:

velocity = momentum * velocity + gradient
param    = param * (1 - lr * weightDecay) - lr * velocity

Weight decay is applied as a tensor-op (MulScalar) — no CPU loop. Momentum helps accelerate SGD in relevant directions and dampens oscillations.

All parameter updates are performed as pure tensor operations — no data is transferred to CPU during Step(). This keeps the hot path entirely on the accelerator when using a GPU backend.

Example:

optimizer := optim.NewSGD(model.Parameters(), optim.SGDConfig{
    LR:          0.01,
    Momentum:    0.9,
    WeightDecay: 1e-4,
})

for epoch := range epochs {
    loss := train_step(model, batch)
    grads := autodiff.Backward(loss, backend)
    optimizer.Step(grads)
    optimizer.ZeroGrad()
}

func NewSGD[B tensor.Backend](params []*nn.Parameter[B], config SGDConfig, backend B) *SGD[B]

sgd := optim.NewSGD(model.Parameters(), optim.SGDConfig{
    LR:       0.01,
    Momentum: 0.9,
})

func (s *SGD[B]) GetLR() float32

func (s *SGD[B]) LoadStateDict(stateDict map[string]*tensor.RawTensor) error

func (s *SGD[B]) SetLR(lr float32)

func (s *SGD[B]) StateDict() map[string]*tensor.RawTensor

func (s *SGD[B]) Step(grads map[*tensor.RawTensor]*tensor.RawTensor)

func (s *SGD[B]) ZeroGrad()

type SGDConfig ¶

type SGDConfig struct {
	LR          float32 // Learning rate (default: 0.01)
	Momentum    float32 // Momentum factor (default: 0.0, range: [0, 1))
	WeightDecay float32 // L2 weight-decay coefficient (default: 0.0)
}

SGDConfig holds configuration for SGD optimizer.