heavyball 0.7.0

Efficient optimizers

HeavyBall

A simple package of efficient optimizers

The goal is not to thrive for completeness, full maintenance or abstraction, but instead to provide a simple largely static alternative to torch.optim with more and better optimizers.

Features

• Stochastic Rounding: FP32 convergence with BF16 parameters
• Inplace EMA: Same math, but less memory, less compute and higher stability
• Foreach: Fast multi-tensor application
• PaLM Beta2: Fast initial convergence, stable late convergence
• ScheduleFree: No learning rate schedule, but better convergence

Getting started

pip install heavyball

import torch
import heavyball

# Create a model
model = torch.nn.Linear(16, 1)

# Create an optimizer
optimizer = heavyball.PaLMForeachSFAdamW(model.parameters(), lr=1e-3)

x = torch.randn(128, 16)
y = torch.randn(128, 1)

for _ in range(1000):
    optimizer.zero_grad()
    loss = torch.nn.functional.mse_loss(model(x), y)
    loss.backward()
    optimizer.step()

Optimizers

Name	Description	Advantages / Disadvantages
ForeachAdamW	More efficient (speed, memory) AdamW	+ Faster than AdamW + Possibly more (numerically) stable
ForeachLaProp	More efficient (speed, memory) LaProp	+ Same cost as AdamW + Marginally better converence (better proofs) + Higher hyperparameter stability - Not a guaranteed win (can be neutral) - No "Slingshot"
ForeachADOPT	More efficient (speed, memory) ADOPT	+ Same cost as AdamW + Rigorous mathematical convergence proofs, even for challenging models (GANs) - Empirically underperforms LaProp - no bf16
ForeachSFAdamW	More efficient (speed, memory) ScheduleFree AdamW	+ Same cost as AdamW, but better eval perf + Full control over hyperparameters
PaLMForeachSFAdamW	ForeachSFAdamW with PaLM's beta2 schedule	+ Same cost as AdamW, but better eval perf + Less control, but faster early and more stable late convergence + ScheduleFree - slow early convergence
ForeachSOAP	More efficient (speed, memory) SOAP	+ Fastest convergence (loss-at-step) + Full control over hyperparameters - more memory usage - more hyperparameters - higher overhead than AdamW (can be ammortized; better loss-at-second)
PaLMForeachSOAP	ForeachSOAP with PaLM's beta2 schedule	+ Fastest convergence (loss-at-step) + Less control, but faster early and more stable late convergence - more memory usage - more hyperparameters - higher overhead than AdamW (can be ammortized; better loss-at-second)
SFPaLMForeachSOAP	ScheduleFree PaLMForeachSOAP	+ Fast convergence (loss-at-step) + less memory usage than PaLMForeachSOAP (more tham AdamW) - slower initial convergence than PaLMForeachSOAP (but allows higher LRs) - higher overhead than AdamW (can be ammortized)
PrecondScheduleSFPaLMForeachSOAP	SFPaLMForeachSOAP with preconditioner schedule, matching the error of PrecondEvery=2 with the cost of PrecondEvery=512	+ Better initial convergence than SFPaLMForeachSOAP + Significantly faster (sec/it) later + less memory usage than PaLMForeachSOAP (more tham AdamW) - slower initial convergence than PaLMForeachSOAP (but allows higher LRs) - higher overhead than AdamW (can be ammortized)