Adam Optimizer (Adaptive Moment Estimation)

Adam optimizer with adaptive learning rates per parameter. Combines ideas from RMSprop and momentum.

Algorithm

m_t = β₁ * m_{t-1} + (1 - β₁) * g_t          # First moment (momentum)
v_t = β₂ * v_{t-1} + (1 - β₂) * g_t²         # Second moment (uncentered variance)
m̂_t = m_t / (1 - β₁^t)                        # Bias correction
v̂_t = v_t / (1 - β₂^t)                        # Bias correction
θ_t = θ_{t-1} - α * m̂_t / (√v̂_t + ε)       # Parameter update

Typical Hyperparameters

• Learning rate (α): 0.001
• β₁ (momentum): 0.9
• β₂ (RMS decay): 0.999
• ε (numerical stability): 1e-8

Usage Example

-- Create optimizer
let opt := AdamConfig.default.withLearningRate 0.001

-- Optimization step
let (newParams, newState) := opt.step params grads state

References

• Kingma & Ba (2014): "Adam: A Method for Stochastic Optimization"
• https://arxiv.org/abs/1412.6980

structure Hesper.Optimizer.Adam.AdamConfig : Type

Adam optimizer configuration

• learningRate : Float

  Learning rate (typical: 0.001)

• beta1 : Float

  Exponential decay rate for first moment estimates (typical: 0.9)

• beta2 : Float

  Exponential decay rate for second moment estimates (typical: 0.999)

• epsilon : Float

  Small constant for numerical stability (typical: 1e-8)

• weightDecay : Float

  Weight decay (L2 regularization)

• amsgrad : Bool

  Use AMSGrad variant (maintains max of v_t)

instance Hesper.Optimizer.Adam.instInhabitedAdamConfig : Inhabited AdamConfig

Equations

• Hesper.Optimizer.Adam.instInhabitedAdamConfig = { default := Hesper.Optimizer.Adam.instInhabitedAdamConfig.default }

def Hesper.Optimizer.Adam.instInhabitedAdamConfig.default : AdamConfig

Equations

• Hesper.Optimizer.Adam.instInhabitedAdamConfig.default = { learningRate := default, beta1 := default, beta2 := default, epsilon := default, weightDecay := default, amsgrad := default }

def Hesper.Optimizer.Adam.instReprAdamConfig.repr : AdamConfig → Nat → Std.Format

Equations

• One or more equations did not get rendered due to their size.

instance Hesper.Optimizer.Adam.instReprAdamConfig : Repr AdamConfig

Equations

• Hesper.Optimizer.Adam.instReprAdamConfig = { reprPrec := Hesper.Optimizer.Adam.instReprAdamConfig.repr }

def Hesper.Optimizer.Adam.AdamConfig.default : AdamConfig

Default Adam configuration (recommended hyperparameters)

Equations

• Hesper.Optimizer.Adam.AdamConfig.default = { }

def Hesper.Optimizer.Adam.AdamConfig.withLearningRate (config : AdamConfig) (lr : Float) : AdamConfig

Create Adam with specified learning rate

Equations

• config.withLearningRate lr = { learningRate := lr, beta1 := config.beta1, beta2 := config.beta2, epsilon := config.epsilon, weightDecay := config.weightDecay, amsgrad := config.amsgrad }

def Hesper.Optimizer.Adam.AdamConfig.withWeightDecay (config : AdamConfig) (wd : Float) : AdamConfig

Create Adam with weight decay (AdamW variant)

Equations

• config.withWeightDecay wd = { learningRate := config.learningRate, beta1 := config.beta1, beta2 := config.beta2, epsilon := config.epsilon, weightDecay := wd, amsgrad := config.amsgrad }

def Hesper.Optimizer.Adam.AdamConfig.withBetas (config : AdamConfig) (beta1 beta2 : Float) : AdamConfig

Create Adam with custom beta parameters

Equations

• config.withBetas beta1 beta2 = { learningRate := config.learningRate, beta1 := beta1, beta2 := beta2, epsilon := config.epsilon, weightDecay := config.weightDecay, amsgrad := config.amsgrad }

def Hesper.Optimizer.Adam.AdamConfig.withAMSGrad (config : AdamConfig) : AdamConfig

Enable AMSGrad variant

Equations

• config.withAMSGrad = { learningRate := config.learningRate, beta1 := config.beta1, beta2 := config.beta2, epsilon := config.epsilon, weightDecay := config.weightDecay, amsgrad := true }

structure Hesper.Optimizer.Adam.AdamState : Type

Adam optimizer state (moment estimates)

• m : Array (Array Float)

  First moment estimates (momentum) for each parameter

• v : Array (Array Float)

  Second moment estimates (uncentered variance) for each parameter

• vMax : Array (Array Float)

  Maximum of v_t (for AMSGrad)

• step : Nat

  Number of steps taken

instance Hesper.Optimizer.Adam.instInhabitedAdamState : Inhabited AdamState

Equations

• Hesper.Optimizer.Adam.instInhabitedAdamState = { default := Hesper.Optimizer.Adam.instInhabitedAdamState.default }

def Hesper.Optimizer.Adam.instInhabitedAdamState.default : AdamState

Equations

• Hesper.Optimizer.Adam.instInhabitedAdamState.default = { m := default, v := default, vMax := default, step := default }

def Hesper.Optimizer.Adam.instReprAdamState.repr : AdamState → Nat → Std.Format

Equations

• One or more equations did not get rendered due to their size.

instance Hesper.Optimizer.Adam.instReprAdamState : Repr AdamState

Equations

• Hesper.Optimizer.Adam.instReprAdamState = { reprPrec := Hesper.Optimizer.Adam.instReprAdamState.repr }

def Hesper.Optimizer.Adam.AdamState.init (paramShapes : Array Nat) : AdamState

Initialize Adam state for given parameter shapes

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Optimizer.Adam.AdamState.fromParams (params : Array (Array Float)) : AdamState

Initialize Adam state from parameters

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Optimizer.Adam.updateParam (config : AdamConfig) (param grad m v vMax : Array Float) (step : Nat) : Array Float × Array Float × Array Float × Array Float

Update a single parameter with Adam

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Optimizer.Adam.step (config : AdamConfig) (params grads : Array (Array Float)) (state : AdamState) : Array (Array Float) × AdamState

Perform one optimization step for all parameters

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Optimizer.Adam.effectiveLearningRate (config : AdamConfig) (step : Nat) : Float

Helper: Compute effective learning rate at step t

Equations

• Hesper.Optimizer.Adam.effectiveLearningRate config step = config.learningRate * (1.0 - config.beta2.pow (Float.ofNat (step + 1))).sqrt / (1.0 - config.beta1.pow (Float.ofNat (step + 1)))

def Hesper.Optimizer.Adam.momentumNorm (state : AdamState) : Float

Helper: Compute L2 norm of first moments (momentum magnitude)

Equations

• Hesper.Optimizer.Adam.momentumNorm state = (Array.foldl (fun (acc : Float) (m : Array Float) => acc + Array.foldl (fun (acc2 mi : Float) => acc2 + mi * mi) 0.0 m) 0.0 state.m).sqrt

def Hesper.Optimizer.Adam.varianceNorm (state : AdamState) : Float

Helper: Compute L2 norm of second moments (variance magnitude)

Equations

• Hesper.Optimizer.Adam.varianceNorm state = (Array.foldl (fun (acc : Float) (v : Array Float) => acc + Array.foldl (fun (acc2 vi : Float) => acc2 + vi * vi) 0.0 v) 0.0 state.v).sqrt

structure Hesper.Optimizer.Adam.AdamStats : Type

Helper: Get statistics for monitoring

• step : Nat
• effectiveLR : Float
• momentumNorm : Float
• varianceNorm : Float

def Hesper.Optimizer.Adam.instReprAdamStats.repr : AdamStats → Nat → Std.Format

Equations

• One or more equations did not get rendered due to their size.

instance Hesper.Optimizer.Adam.instReprAdamStats : Repr AdamStats

Equations

• Hesper.Optimizer.Adam.instReprAdamStats = { reprPrec := Hesper.Optimizer.Adam.instReprAdamStats.repr }

def Hesper.Optimizer.Adam.getStats (config : AdamConfig) (state : AdamState) : AdamStats

Equations

• One or more equations did not get rendered due to their size.