Advanced Optimizers & Schedulers

1/**2 * Example 41: Advanced Optimizers & Schedulers3 *4 * New in v1.0.0: RAdam, LAMB, LARS optimizers and CyclicLR,5 * CosineAnnealingWarmRestarts, PolynomialLR, LambdaLR, SequentialLR schedulers.6 */78import { GradTensor, type Tensor, tensor } from "deepbox/ndarray";9import { Linear, mseLoss, ReLU, Sequential } from "deepbox/nn";10import {11  Adam,12  CosineAnnealingWarmRestarts,13  CyclicLR,14  LAMB,15  LARS,16  LambdaLR,17  PolynomialLR,18  RAdam,19  SequentialLR,20  StepLR,21} from "deepbox/optim";2223console.log("=".repeat(60));24console.log("Example 41: Advanced Optimizers & Schedulers");25console.log("=".repeat(60));2627// ============================================================================28// Helper: simple training demo29// ============================================================================3031function trainDemo(32  model: Sequential,33  optimizerName: string,34  optimizer: { step: () => void; zeroGrad: () => void; lr: number },35  epochs = 1036): void {37  const xTrain = tensor([38    [1, 2],39    [3, 4],40    [5, 6],41    [7, 8],42  ]);43  const yTrain = tensor([[3], [7], [11], [15]]);4445  console.log(`\n  ${optimizerName} (initial lr=${optimizer.lr.toFixed(6)}):`);4647  for (let epoch = 1; epoch <= epochs; epoch++) {48    optimizer.zeroGrad();49    const pred = model.forward(xTrain);50    const loss = mseLoss(pred, yTrain);51    const scalarLoss: Tensor = GradTensor.isGradTensor(loss) ? loss.tensor : loss;52    const lossVal = Number(scalarLoss.data[scalarLoss.offset]);5354    if (GradTensor.isGradTensor(loss)) {55      loss.backward();56    }57    optimizer.step();5859    if (epoch === 1 || epoch === epochs || epoch % 5 === 0) {60      console.log(61        `    Epoch ${String(epoch).padStart(3)}: loss=${lossVal.toFixed(6)}, lr=${optimizer.lr.toFixed(6)}`62      );63    }64  }65}6667// ============================================================================68// Part 1: RAdam (Rectified Adam)69// ============================================================================70console.log("\n🚀 Part 1: RAdam (Rectified Adam)");71console.log("-".repeat(60));7273// RAdam auto-adjusts the adaptive learning rate based on variance of gradients74// No need for learning rate warmup75const model1 = new Sequential(new Linear(2, 8), new ReLU(), new Linear(8, 1));76const radam = new RAdam(model1.parameters(), { lr: 0.01 });7778console.log("RAdam: Rectified Adam — no warmup needed");79console.log("  Auto-adjusts adaptive LR based on gradient variance");80trainDemo(model1, "RAdam", radam);8182// ============================================================================83// Part 2: LAMB (Layer-wise Adaptive Moments)84// ============================================================================85console.log("\n🐑 Part 2: LAMB");86console.log("-".repeat(60));8788// LAMB scales gradients layer-wise — great for large batch training89const model2 = new Sequential(new Linear(2, 8), new ReLU(), new Linear(8, 1));90const lamb = new LAMB(model2.parameters(), { lr: 0.01 });9192console.log("LAMB: Layer-wise Adaptive Moments — ideal for large batch training");93trainDemo(model2, "LAMB", lamb);9495// ============================================================================96// Part 3: LARS (Layer-wise Adaptive Rate Scaling)97// ============================================================================98console.log("\n🏔️  Part 3: LARS");99console.log("-".repeat(60));100101// LARS adjusts learning rate per layer based on weight/gradient norms102const model3 = new Sequential(new Linear(2, 8), new ReLU(), new Linear(8, 1));103const lars = new LARS(model3.parameters(), { lr: 0.01 });104105console.log("LARS: Layer-wise Adaptive Rate Scaling — for very large batches");106trainDemo(model3, "LARS", lars);107108// ============================================================================109// Part 4: CyclicLR Scheduler110// ============================================================================111console.log("\n🔄 Part 4: CyclicLR Scheduler");112console.log("-".repeat(60));113114// CyclicLR cycles the learning rate between a base and max value115const model4 = new Sequential(new Linear(2, 8), new ReLU(), new Linear(8, 1));116const adam4 = new Adam(model4.parameters(), { lr: 0.001 });117const cyclicLr = new CyclicLR(adam4, {118  baseLr: 0.001,119  maxLr: 0.01,120  stepSizeUp: 5,121  mode: "triangular",122});123124console.log("CyclicLR: triangular cycling between 0.001 and 0.01");125console.log("  LR schedule over 20 steps:");126for (let i = 0; i < 20; i++) {127  const lrs = cyclicLr.getLr();128  if (i % 4 === 0 || i === 19) {129    console.log(`    Step ${String(i + 1).padStart(3)}: lr=${lrs[0]?.toFixed(6)}`);130  }131  cyclicLr.step();132}133134// ============================================================================135// Part 5: CosineAnnealingWarmRestarts136// ============================================================================137console.log("\n🌊 Part 5: CosineAnnealingWarmRestarts");138console.log("-".repeat(60));139140// Cosine annealing with periodic warm restarts141const model5 = new Sequential(new Linear(2, 8), new ReLU(), new Linear(8, 1));142const adam5 = new Adam(model5.parameters(), { lr: 0.01 });143const cosineWR = new CosineAnnealingWarmRestarts(adam5, {144  T_0: 5, // restart every 5 epochs145  T_mult: 2, // double the period after each restart146  etaMin: 0.001,147});148149console.log("CosineAnnealingWarmRestarts: T_0=5, T_mult=2, etaMin=0.001");150console.log("  LR schedule (restarts at epoch 5, then 15, ...):");151for (let i = 0; i < 20; i++) {152  const lrs = cosineWR.getLr();153  if (i % 3 === 0 || i === 19) {154    console.log(`    Epoch ${String(i + 1).padStart(3)}: lr=${lrs[0]?.toFixed(6)}`);155  }156  cosineWR.step();157}158159// ============================================================================160// Part 6: PolynomialLR161// ============================================================================162console.log("\n📉 Part 6: PolynomialLR");163console.log("-".repeat(60));164165// Polynomial decay from initial LR to end LR166const model6 = new Sequential(new Linear(2, 8), new ReLU(), new Linear(8, 1));167const adam6 = new Adam(model6.parameters(), { lr: 0.01 });168const polyLr = new PolynomialLR(adam6, {169  totalIters: 20,170  power: 2.0,171});172173console.log("PolynomialLR: power=2.0 decay over 20 iterations");174console.log("  LR schedule:");175for (let i = 0; i < 20; i++) {176  const lrs = polyLr.getLr();177  if (i % 4 === 0 || i === 19) {178    console.log(`    Step ${String(i + 1).padStart(3)}: lr=${lrs[0]?.toFixed(6)}`);179  }180  polyLr.step();181}182183// ============================================================================184// Part 7: LambdaLR185// ============================================================================186console.log("\n🔧 Part 7: LambdaLR");187console.log("-".repeat(60));188189// LambdaLR uses a custom function to compute the LR multiplier190const model7 = new Sequential(new Linear(2, 8), new ReLU(), new Linear(8, 1));191const adam7 = new Adam(model7.parameters(), { lr: 0.01 });192const lambdaLr = new LambdaLR(adam7, {193  lrLambda: (epoch: number) => 0.95 ** epoch, // exponential decay194});195196console.log("LambdaLR: lr_lambda = 0.95^epoch (exponential decay)");197console.log("  LR schedule:");198for (let i = 0; i < 20; i++) {199  const lrs = lambdaLr.getLr();200  if (i % 4 === 0 || i === 19) {201    console.log(`    Epoch ${String(i + 1).padStart(3)}: lr=${lrs[0]?.toFixed(6)}`);202  }203  lambdaLr.step();204}205206// ============================================================================207// Part 8: SequentialLR208// ============================================================================209console.log("\n📋 Part 8: SequentialLR");210console.log("-".repeat(60));211212// SequentialLR chains multiple schedulers at specified milestones213const model8 = new Sequential(new Linear(2, 8), new ReLU(), new Linear(8, 1));214const adam8 = new Adam(model8.parameters(), { lr: 0.01 });215216// Phase 1: Warmup with LambdaLR (epochs 0-4)217const warmup = new LambdaLR(adam8, {218  lrLambda: (epoch: number) => Math.min(1.0, (epoch + 1) / 5),219});220221// Phase 2: StepLR decay (epochs 5+)222const decay = new StepLR(adam8, { stepSize: 3, gamma: 0.5 });223224const seqLr = new SequentialLR(adam8, {225  schedulers: [warmup, decay],226  milestones: [5],227});228229console.log("SequentialLR: LambdaLR warmup (0-4) → StepLR decay (5+)");230console.log("  LR schedule:");231for (let i = 0; i < 20; i++) {232  const lrs = seqLr.getLr();233  if (i % 3 === 0 || i === 19) {234    console.log(`    Epoch ${String(i + 1).padStart(3)}: lr=${lrs[0]?.toFixed(6)}`);235  }236  seqLr.step();237}238239// ============================================================================240// Part 9: Optimizer Comparison241// ============================================================================242console.log("\n📊 Part 9: Optimizer Comparison");243console.log("-".repeat(60));244245console.log("┌─────────────────────────────┬──────────────────────────────────────┐");246console.log("│ Optimizer                   │ Best For                             │");247console.log("├─────────────────────────────┼──────────────────────────────────────┤");248console.log("│ Adam                        │ General purpose, default choice      │");249console.log("│ RAdam                       │ No warmup needed, stable convergence │");250console.log("│ LAMB                        │ Large batch distributed training     │");251console.log("│ LARS                        │ Very large batch SGD-style training  │");252console.log("└─────────────────────────────┴──────────────────────────────────────┘");253254console.log("\n┌─────────────────────────────┬──────────────────────────────────────┐");255console.log("│ Scheduler                   │ Strategy                             │");256console.log("├─────────────────────────────┼──────────────────────────────────────┤");257console.log("│ CyclicLR                    │ Triangular LR cycling                │");258console.log("│ CosineAnnealingWarmRestarts │ Cosine decay with periodic restarts  │");259console.log("│ PolynomialLR                │ Polynomial decay to end LR           │");260console.log("│ LambdaLR                    │ Custom function-based scheduling     │");261console.log("│ SequentialLR                │ Chain multiple schedulers at epochs   │");262console.log("└─────────────────────────────┴──────────────────────────────────────┘");263264// ============================================================================265// Summary266// ============================================================================267console.log("\n💡 Key Takeaways");268console.log("-".repeat(60));269console.log("• RAdam: no warmup needed, automatically adjusts adaptive learning rate");270console.log("• LAMB: layer-wise scaling for large batch training (keeps per-layer LR)");271console.log("• LARS: layer-wise rate scaling for SGD-style very large batch training");272console.log("• CyclicLR: avoids local minima by cycling LR between base and max");273console.log("• CosineWarmRestarts: periodic warm restarts explore new loss basins");274console.log("• PolynomialLR: smooth polynomial decay over fixed iterations");275console.log("• LambdaLR: fully custom scheduling via user-defined functions");276console.log("• SequentialLR: combine warmup + decay phases at milestone epochs");277278console.log("\n✅ Advanced Optimizers & Schedulers Example Complete!");279console.log("=".repeat(60));280