Efficient Hyperparameter Optimization for Large Models

# Illustrative sketch (not from the paper)
import torch
import numpy as np

# 1. Small‑scale proxy runs produce (token_budget, lr_opt) pairs
# token_budgets and lr_optimal are placeholders for actual collected data
token_budgets = torch.tensor([1e6, 5e6, 1e7])  # example token counts
lr_optimal = torch.tensor([1e-3, 5e-4, 2.5e-4])  # learned optimal LR per budget

# 2. Fit a second‑order polynomial η(t) = a*t^2 + b*t + c on the proxy data
coeffs = np.polyfit(token_budgets.numpy(), lr_optimal.numpy(), deg=2)
# coeffs = [a, b, c]

# 3. Predict optimal LR for a target massive token budget (e.g., 1e13)
target_budget = 1e13
eta_star = np.polyval(coeffs, target_budget)

# 4. Log‑log linear regression on (budget, η*) pairs to extrapolate further
log_budget = torch.log(token_budgets)
log_eta = torch.log(lr_optimal)
# Simple linear regression: log_eta = m * log_budget + c
X = torch.stack([log_budget, torch.ones_like(log_budget)], dim=1)
# Solve for m, c via least squares
m, c = torch.linalg.lstsq(X, log_eta.unsqueeze(1)).solution.squeeze()
# Extrapolate to the massive horizon using the fitted line
eta_extrap = torch.exp(m * torch.log(torch.tensor(target_budget)) + c)

# 5. EMA checkpoint extraction (illustrative, no decay schedule)
ema_decay = 0.999
ema_state = None
for step, model_state in enumerate(proxy_training_loop()):
    if ema_state is None:
        ema_state = model_state.clone()
    else:
        ema_state = ema_decay * ema_state + (1 - ema_decay) * model_state
# ema_state now holds the EMA checkpoint

# 6. Use Muon optimizer (placeholder – actual implementation resides in the Muon library)
# optimizer = Muon(model.parameters(), lr=eta_extrap)

// Illustrative sketch (not from the paper)
const tf = require('@tensorflow/tfjs-node');

// 1. Proxy data: token budgets and optimal learning rates (placeholders)
const tokenBudgets = tf.tensor1d([1e6, 5e6, 1e7]); // example budgets
const lrOptimal = tf.tensor1d([1e-3, 5e-4, 2.5e-4]); // optimal LR per budget

// 2. Fit a quadratic polynomial η(t) = a*t^2 + b*t + c using least‑squares
function polyFit(x, y, degree) {
  const X = [];
  for (let i = 0; i < x.shape[0]; i++) {
    const xi = x.arraySync()[i];
    const row = [];
    for (let d = degree; d >= 0; d--) row.push(Math.pow(xi, d));
    X.push(row);
  }
  const XT = tf.tensor2d(X).transpose();
  const coeffs = tf.linalg.pinv(XT.matMul(tf.tensor2d(X))).matMul(XT).matMul(y.reshape([-1, 1]));
  return coeffs.squeeze().arraySync(); // [a, b, c]
}
const coeffs = polyFit(tokenBudgets, lrOptimal, 2);

// 3. Predict optimal LR for a massive token budget (e.g., 1e13)
const targetBudget = 1e13;
const etaStar = coeffs[0] * Math.pow(targetBudget, 2) + coeffs[1] * targetBudget + coeffs[2];

// 4. Log‑log linear regression to extrapolate further
const logBudget = tokenBudgets.log();
const logEta = lrOptimal.log();
const ones = tf.onesLike(logBudget);
const Xlog = tf.stack([logBudget, ones], 1);
const solution = tf.linalg.lstsq(Xlog, logEta.reshape([-1, 1]));
const m = solution.squeeze().arraySync()[0];
const c = solution.squeeze().arraySync()[1];
const etaExtrap = Math.exp(m * Math.log(targetBudget) + c);

// 5. EMA checkpoint extraction (illustrative, no decay schedule)
let emaState = null;
const emaDecay = 0.999;
for (const modelState of proxyTrainingLoop()) {
  if (!emaState) {
    emaState = modelState.clone();
  } else {
    emaState = emaState.mul(emaDecay).add(modelState.mul(1 - emaDecay));
  }
}
// emaState now holds the EMA checkpoint

// 6. Muon optimizer placeholder – actual implementation resides in the Muon library
// const optimizer = new MuonOptimizer(model.parameters(), { lr: etaExtrap });