diff --git a/docs/OPTIMIZER_API_VERIFICATION.md b/docs/OPTIMIZER_API_VERIFICATION.md index 3ac62814..9907d845 100644 --- a/docs/OPTIMIZER_API_VERIFICATION.md +++ b/docs/OPTIMIZER_API_VERIFICATION.md @@ -141,12 +141,12 @@ The only remaining work for complete optimizer functionality: 2. **Additional optimizer classes**: - ✅ Adam (structure and validation complete, blocked on missing core ops) + - ✅ Adamax (fully implemented with all required ops) + - ✅ Lion (fully implemented with all required ops) + - ✅ RMSprop (fully implemented with all required ops) - AdamW - - RMSprop - Adagrad - AdaDelta - - Adamax - - Lion - Adafactor - Muon diff --git a/node/CHECKLIST.md b/node/CHECKLIST.md index 4f7848e5..c4d6f323 100644 --- a/node/CHECKLIST.md +++ b/node/CHECKLIST.md @@ -46,6 +46,7 @@ Reference source map: see `PORTING_PLAN.md` for the list of upstream C++ and Pyt - [ ] Complete `SGD.applySingle()` implementation (requires subtract, scalar ops, astype). - [x] Add `Adam` optimizer (structure and validation complete, blocked on missing core ops). - [ ] Complete `Adam.applySingle()` implementation (requires subtract, divide, square, sqrt, rsqrt, power). +- [x] Add `Adamax` optimizer (fully implemented with all required ops). - [x] Add `Lion` optimizer (fully implemented with all required ops). - [x] Add `RMSprop` optimizer (fully implemented with all required ops). - [ ] Add `AdamW` optimizer. diff --git a/node/src/optimizers/index.ts b/node/src/optimizers/index.ts index 28697520..9b42b69b 100644 --- a/node/src/optimizers/index.ts +++ b/node/src/optimizers/index.ts @@ -6,7 +6,7 @@ */ import MLXArray, { zeros, zerosLike } from '../core/array'; -import { add, multiply, subtract, sign, square, sqrt, divide, matmul, reshape, transpose } from '../core/ops'; +import { add, multiply, subtract, sign, square, sqrt, divide, abs, maximum, matmul, reshape, transpose } from '../core/ops'; import { treeMap } from '../utils'; /** @@ -453,6 +453,92 @@ export class Adam extends Optimizer { } } +/** + * Adamax Optimizer Options + */ +export interface AdamaxOptions { + /** The learning rate */ + learningRate: SchedulableParam; + /** The coefficients (β₁, β₂) used for computing running averages of gradient and the exponentially weighted infinity norm (default: [0.9, 0.999]) */ + betas?: [number, number]; + /** The term ε added to the denominator to improve numerical stability (default: 1e-8) */ + eps?: number; +} + +/** + * The Adamax optimizer, a variant of Adam based on the infinity norm. + * + * Implements the Adamax algorithm from "Adam: A Method for Stochastic Optimization" (Kingma & Ba, 2015). + * + * The algorithm updates parameters as follows: + * + * m_{t+1} = β₁ * m_t + (1 - β₁) * g_t + * v_{t+1} = max(β₂ * v_t, |g_t|) + * w_{t+1} = w_t - λ * m_{t+1} / (v_{t+1} + ε) + * + * where λ is the learning rate, m_t is the first moment estimate, + * v_t is the exponentially weighted infinity norm, g_t is the gradient, + * and ε is a small constant for numerical stability. + * + * @example + * ```typescript + * const optimizer = new Adamax({ learningRate: 0.002 }); + * // ... during training: + * // const updatedParams = optimizer.applyGradients(gradients, parameters); + * ``` + */ +export class Adamax extends Adam { + constructor(options: AdamaxOptions) { + const { + learningRate, + betas = [0.9, 0.999], + eps = 1e-8 + } = options; + + if (eps < 0.0) { + throw new Error(`Epsilon value should be >=0, ${eps} was provided instead`); + } + + // Call parent constructor with biasCorrection disabled + // Following the original paper, Adamax omits bias correction + super({ learningRate, betas, eps, biasCorrection: false }); + } + + protected initSingle(parameter: MLXArray, state: Record): void { + // Initialize first moment (m) and exponentially weighted infinity norm (v) with zerosLike + state.m = zerosLike(parameter); + state.v = zerosLike(parameter); + } + + protected applySingle( + gradient: MLXArray, + parameter: MLXArray, + state: Record + ): MLXArray { + const lr = this.learningRate; + const [b1, b2] = this.betas; + const eps = this.eps; + + // Get moments from state + let m = state.m; + let v = state.v; + + // Update biased first moment estimate: m = β₁ * m + (1 - β₁) * g + m = add(multiply(b1, m), multiply(1 - b1, gradient)); + + // Update exponentially weighted infinity norm: v = max(β₂ * v, |g|) + v = maximum(multiply(b2, v), abs(gradient)); + + // Store updated moments + state.m = m; + state.v = v; + + // Compute update: w = w - lr * m / (v + ε) + const update = divide(multiply(lr, m), add(v, eps)); + return subtract(parameter, update); + } +} + export interface LionOptions { learningRate: SchedulableParam; betas?: [number, number]; @@ -865,6 +951,7 @@ export default { Optimizer, SGD, Adam, + Adamax, Lion, Adagrad, RMSprop, diff --git a/node/test/optimizers.test.ts b/node/test/optimizers.test.ts index acb9ff6d..a0df4ba1 100644 --- a/node/test/optimizers.test.ts +++ b/node/test/optimizers.test.ts @@ -1,6 +1,6 @@ import { describe, it } from 'node:test'; import assert from 'node:assert'; -import { SGD, Adam, Lion, Adagrad, RMSprop, Muon, Optimizer } from '../src/optimizers'; +import { SGD, Adam, Adamax, Lion, Adagrad, RMSprop, Muon, Optimizer } from '../src/optimizers'; import { zeros } from '../src/core/array'; describe('mlx.optimizers', () => { @@ -212,6 +212,89 @@ describe('mlx.optimizers', () => { // These should be added once those operations are available }); + describe('mlx.optimizers.Adamax', () => { + it('should create Adamax optimizer with learning rate', () => { + const optimizer = new Adamax({ learningRate: 0.002 }); + assert.ok(optimizer instanceof Optimizer); + assert.ok(optimizer instanceof Adam); + assert.ok(optimizer instanceof Adamax); + assert.deepStrictEqual(optimizer.betas, [0.9, 0.999]); + assert.strictEqual(optimizer.eps, 1e-8); + assert.strictEqual(optimizer.biasCorrection, false); + }); + + it('should create Adamax optimizer with custom betas', () => { + const optimizer = new Adamax({ + learningRate: 0.002, + betas: [0.95, 0.9999] + }); + assert.deepStrictEqual(optimizer.betas, [0.95, 0.9999]); + }); + + it('should create Adamax optimizer with custom epsilon', () => { + const optimizer = new Adamax({ + learningRate: 0.002, + eps: 1e-7 + }); + assert.strictEqual(optimizer.eps, 1e-7); + }); + + it('should throw error for negative epsilon', () => { + assert.throws( + () => new Adamax({ learningRate: 0.002, eps: -1e-8 }), + /Epsilon value should be >=0/ + ); + }); + + it('should have learning rate in state', () => { + const optimizer = new Adamax({ learningRate: 0.002 }); + const lr = optimizer.learningRate; + assert.ok(lr); + // Check that it's an MLXArray + assert.ok(lr.toTypedArray !== undefined); + }); + + it('should initialize state for parameters', () => { + const optimizer = new Adamax({ learningRate: 0.002 }); + const params = { + weight: zeros([3]), + bias: zeros([1]) + }; + + optimizer.init(params); + + // Check that state was initialized + assert.ok(optimizer.state); + assert.ok('step' in optimizer.state); + assert.ok('learning_rate' in optimizer.state); + assert.ok('weight' in optimizer.state); + assert.ok('bias' in optimizer.state); + + // Check that moment estimates were initialized + assert.ok('m' in optimizer.state.weight); + assert.ok('v' in optimizer.state.weight); + assert.ok('m' in optimizer.state.bias); + assert.ok('v' in optimizer.state.bias); + }); + + it('should track step count', () => { + const optimizer = new Adamax({ learningRate: 0.002 }); + const step = optimizer.step; + assert.ok(step); + assert.strictEqual(step.toTypedArray()[0], 0); + }); + + it('should allow setting learning rate', () => { + const optimizer = new Adamax({ learningRate: 0.002 }); + optimizer.learningRate = 0.0001; + // Note: This test is placeholder since we can't properly set scalar values yet + assert.ok(optimizer.learningRate); + }); + + // Note: Full integration tests with applyGradients would be added here + // once we have all operations working properly + }); + describe('Lion', () => { it('should create Lion optimizer with learning rate', () => { const optimizer = new Lion({ learningRate: 0.0001 });