afor-optimizer

AFOR

AFOR is a tensor-wise adaptive forgetting optimizer for PyTorch. It adapts the second-moment decay coefficient from local gradient residuals, direction consistency, and online Z-score normalization.

Installation

pip install afor-optimizer

Usage

import torch
from AFOR import afor

model = torch.nn.Linear(10, 2)
optimizer = afor(
    model.parameters(),
    lr=1e-3,
    betas=(0.9, 0.999),
    beta2_min=0.99,
    dir_weight=1.0,
    weight_decay=1e-4,
)

inputs = torch.randn(16, 10)
targets = torch.randint(0, 2, (16,))
loss = torch.nn.functional.cross_entropy(model(inputs), targets)
loss.backward()
optimizer.step()
optimizer.zero_grad()

Main Parameters

• lr: learning rate.
• betas: first-moment and initial second-moment coefficients.
• beta2_min: lower bound for the adaptive second-moment coefficient.
• dir_weight: weight of gradient-direction consistency.
• eps: numerical stability term.
• weight_decay: decoupled weight decay.
• fast_ref: use the larger fast/slow noise estimate as the noise reference.

AFOR currently supports dense gradients only.

License

Released under the MIT License. See LICENSE.