netcl wiki
api

netcl.optim: Optimizers & Learning Rate Schedulers

netcl.optim: Optimizers & Learning Rate Schedulers

The netcl.optim module provides fused GPU parameter update algorithms and learning rate schedules. It consumes gradients computed by the autograd tape, updates model parameters, clears buffers via zero_grad(), and dynamically adjusts learning rates across epochs.

[!TIP] Theory & Curriculum Link: For the first-principles mathematical derivation of SGD, momentum ravines, adaptive gradient moments, and decoupled weight decay, see Curriculum: Optimization Algorithms.

1. Quick Example: Model Training with AdamW

import numpy as np
import netcl.autograd as ag
import netcl.nn as nn
import netcl.optim as opt
from netcl.core.device import manager
from netcl.core.tensor import Tensor

q = manager.default("auto").queue

# 1. Initialize model and optimizer
model = nn.Linear(q, 10, 2)
optimizer = opt.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)

# 2. Input features and targets on GPU device
x = Tensor.from_host(q, np.random.randn(4, 10).astype(np.float32))
y = Tensor.from_host(q, np.array([0, 1, 0, 1], dtype=np.int32))

# 3. Complete training step
with ag.Tape() as tape:
    logits = model(ag.tensor(x))
    loss = ag.cross_entropy(logits, ag.tensor(y))

tape.backward(loss)
optimizer.step(max_norm=1.0)  # Fused norm clipping + parameter update
optimizer.zero_grad()

print(f"Loss after step: {loss.value.to_host()[0]:.4f}")

2. Optimizer Selection Guide

Class Algorithm Core Hyperparameters Optimal Application
SGD Stochastic Gradient Descent lr, momentum, weight_decay Classical computer vision baselines
Momentum Polyak Heavy-Ball Momentum lr, momentum, weight_decay Standard ConvNets with momentum 0.9
Adam Adaptive Moment Estimation lr, betas, eps, weight_decay General deep learning with coupled L2 regularization
AdamW Decoupled Weight Decay Adam lr, betas, eps, weight_decay Transformers, LLMs, Vision Transformers
RMSProp Root Mean Square Propagation lr, alpha, eps, momentum Recurrent networks, reinforcement learning

3. Mathematical Formulation of Update Rules

SGD with Momentum

For parameter $\theta$ and gradient $\mathbf{g}_t = \nabla_\theta \mathcal{L}$:

$$\mathbf{v}_t = \mu \mathbf{v}_{t-1} + \mathbf{g}_t + \lambda \theta_{t-1}$$

$$\theta_t = \theta_{t-1} - \gamma \mathbf{v}_t$$

Where $\gamma$ is learning rate (lr), $\mu$ is momentum, and $\lambda$ is weight decay.

opt_sgd = opt.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-4)

Adam and AdamW

Maintains exponential moving averages of gradients ($\mathbf{m}_t$) and squared gradients ($\mathbf{v}_t$):

$$\mathbf{m}_t = \beta_1 \mathbf{m}_{t-1} + (1 - \beta_1) \mathbf{g}_t$$

$$\mathbf{v}_t = \beta_2 \mathbf{v}_{t-1} + (1 - \beta_2) \mathbf{g}_t^2$$

Bias corrections:

$$\hat{\mathbf{m}}_t = \frac{\mathbf{m}_t}{1 - \beta_1^t}, \quad \hat{\mathbf{v}}_t = \frac{\mathbf{v}_t}{1 - \beta_2^t}$$

  • Adam (Coupled L2 Regularization):

$$\theta_t = \theta_{t-1} - \gamma \frac{\hat{\mathbf{m}}_t}{\sqrt{\hat{\mathbf{v}}_t} + \epsilon}$$

  • AdamW (Decoupled Weight Decay):

$$\theta_t = \theta_{t-1} - \gamma \left( \frac{\hat{\mathbf{m}}_t}{\sqrt{\hat{\mathbf{v}}_t} + \epsilon} + \lambda \theta_{t-1} \right)$$

opt_adamw = opt.AdamW(model.parameters(), lr=3e-4, betas=(0.9, 0.999), weight_decay=0.01)

Fused Gradient Clipping: step(max_norm=1.0)

Passing max_norm to opt.step() fuses gradient norm calculation, clipping, and parameter updates into a single OpenCL kernel pass, eliminating host roundtrips.

4. Gradient Norm Clipping

Prevents gradient explosions in deep architectures by rescaling the parameter gradient vector to maximum norm $\tau$:

$$\mathbf{g} \leftarrow \mathbf{g} \cdot \min\left(1, \frac{\tau}{\|\mathbf{g}\|_2 + 10^{-6}}\right)$$

from netcl.optim import clip_grad_norm, clip_grad_norm_device

# Device-side fused clipping without CPU host synchronization
clip_grad_norm_device(model.parameters(), max_norm=1.0)

5. Learning Rate Schedulers

Schedulers in netcl compute dynamic learning rates across epochs or iteration steps:

CosineAnnealingLR

Reduces learning rate along a cosine curve from $\eta_{\max}$ down to $\eta_{\min}$ over $T_{\max}$ epochs:

$$\eta_t = \eta_{\min} + \frac{1}{2} (\eta_{\max} - \eta_{\min}) \left( 1 + \cos\left(\frac{t}{T_{\max}} \pi\right) \right)$$

from netcl.optim import CosineAnnealingLR

sched = CosineAnnealingLR(max_lr=3e-4, min_lr=1e-6, T_max=50)

for epoch in range(50):
    optimizer.lr = sched.step()

WarmupCosine

Combines linear warmup over $T_{\text{warmup}}$ steps with subsequent cosine decay:

$$\eta_t = \begin{cases} \eta_{\text{base}} \cdot \frac{t + 1}{T_{\text{warmup}}} & \text{if } t < T_{\text{warmup}} \\ \eta_{\min} + \frac{1}{2} (\eta_{\text{base}} - \eta_{\min}) \left( 1 + \cos\left(\frac{t - T_{\text{warmup}}}{T_{\max} - T_{\text{warmup}}} \pi\right) \right) & \text{if } t \ge T_{\text{warmup}} \end{cases}$$

from netcl.optim import WarmupCosine

sched = WarmupCosine(base_lr=3e-4, max_epochs=50, warmup_epochs=5, min_lr=1e-6)
for epoch in range(50):
    optimizer.lr = sched.lr(epoch)

ReduceLROnPlateau

Reduces learning rate by factor when a monitored metric (e.g. validation loss) fails to improve for patience evaluation cycles:

from netcl.optim import ReduceLROnPlateau

plateau = ReduceLROnPlateau(base_lr=1e-3, factor=0.5, patience=5, min_lr=1e-6)
for epoch in range(100):
    val_loss = 0.42
    optimizer.lr = plateau.step(val_loss)