AADL adds Anderson acceleration to existing PyTorch optimizers. It stores a bounded history of parameter iterates, solves a small least-squares problem, and optionally replaces a normal optimizer update with an extrapolated one.

AADL supports QR and normal-equation Anderson kernels, safeguards, mixed precision, conditioning controls, moving-average smoothing, multiple optimizer parameter groups, and PyTorch-native Post-LocalSGD integration.

Python 3.11 or greater
PyTorch (torch>=2.13) and NumPy (numpy>=2.0)

These minimum versions are enforced by the package metadata and requirements.txt.

The quickest way to get a working environment is the provided helper script, which creates a local virtual environment (.venv), installs the dependencies, and installs AADL in editable mode:

./setup_venv.sh                 # uses python3 by default
PYTHON=python3.11 ./setup_venv.sh   # pick a specific interpreter
./setup_venv.sh --recreate      # delete an existing .venv first

Then activate it:

source .venv/bin/activate

If you prefer to manage your own environment:

python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip setuptools wheel
python -m pip install -r requirements.txt   # core dependencies
python -m pip install -e .                  # install AADL (editable)

The examples/ demos need extra packages (torchvision, pandas, scikit-learn, scikit-image, opencv-python, docopt, pyyaml). Install them with:

python -m pip install -r requirements-examples.txt

python -m unittest discover -s tests -t . -v          # fast suite
RUN_SLOW_TESTS=1 python -m unittest discover -s tests -t . -v   # full suite

The slow suite contains numerical convergence experiments whose results can be sensitive to optimizer and PyTorch version changes. The fast suite contains the API, kernel, safeguard, and distributed-policy regression tests.

AADL has three distinct responsibilities:

• AADL.accelerate wraps a local PyTorch optimizer and computes Anderson candidates. It does not synchronize gradients or model parameters.
• PyTorch DDP/Post-LocalSGD owns gradient communication, process groups, and periodic model averaging.
• AADL's distributed acceptance layer compares globally averaged plain and Anderson branches and reduces only scalar loss statistics.

This separation avoids maintaining a second implementation of DDP, LocalSGD, or FedAvg-style parameter averaging inside AADL.

import torch
import torch.nn
import torch.optim
import AADL

model = torch.nn.Linear(8, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)

AADL.accelerate(
    optimizer,
    acceleration_type="anderson",
    relaxation=0.5,
    wait_iterations=0,
    history_depth=10,
    store_each_nth=1,
    frequency=1,
    reg_acc=1e-8,
    safeguard=True,
)

def closure():
    with torch.enable_grad():
        optimizer.zero_grad()
        loss = loss_fn(model(inputs), targets)
        loss.backward()
    return loss

loss = optimizer.step(closure)

• anderson: QR-factorization kernel and the recommended default.
• anderson_normal_equation: normal-equation kernel, which may be faster but is more sensitive to ill-conditioned histories.
• identity: disables Anderson acceleration; with average=True, it retains only moving-average behavior.

• relaxation: mixing weight in (0, 1].
• wait_iterations: ordinary optimizer steps before acceleration begins.
• history_depth: capacity of the FIFO/ring history.
• store_each_nth: cadence for storing parameter iterates.
• frequency: cadence for attempting Anderson acceleration.
• reg_acc: non-negative Tikhonov regularization.
• safeguard: compare the candidate against the post-optimizer plain step. This requires a closure; without one, the candidate is accepted.
• average: enable stochastic-history moving-average smoothing.
• history_device and compute_device: independently place stored history and the small Anderson solve.
• mixing_dtype: None, a torch.dtype, or a dtype string such as "float32" or "float64".
• equilibrate: unit-scale difference-matrix columns before solving.
• filter_condition: drop oldest columns until the requested condition bound is met; 0 disables filtering.
• refinement_steps: mixed-precision iterative-refinement iterations; 0 disables refinement.

All size and cadence arguments are validated. Calling accelerate twice on the same optimizer raises an error; call AADL.remove_acceleration(optimizer) before changing its configuration.

AADL composes with PyTorch's native Post-LocalSGD hook and model averager:

from torch.distributed.algorithms.ddp_comm_hooks.post_localSGD_hook import (
    PostLocalSGDState,
    post_localSGD_hook,
)
from AADL import HistoryResetPeriodicModelAverager, average_and_accept

state = PostLocalSGDState(
    process_group=None,
    subgroup=None,
    start_localSGD_iter=100,
)
ddp_model.register_comm_hook(state, post_localSGD_hook)

local_optimizer = torch.optim.SGD(ddp_model.parameters(), lr=1e-2)
AADL.accelerate(
    local_optimizer,
    acceleration_type="anderson",
    safeguard=False,  # acceptance is decided globally below
)

averager = HistoryResetPeriodicModelAverager(
    local_optimizer, period=4, warmup_steps=100,
)
# In the training loop, use a closure for the two global loss evaluations:
local_optimizer.step(closure)
average_and_accept(
    local_optimizer,
    averager,
    closure,
    policy="vote",            # or "mean_loss"
    vote_threshold=0.5,
    loss_weight=local_batch_size,
)

average_and_accept returns None between averaging boundaries. At a boundary it returns (accepted, candidate_loss, baseline_loss) and leaves every rank on the same selected global parameters.

Available global policies are:

• vote: accept when at least vote_threshold of ranks report a lower local loss for the shared global Anderson candidate.
• mean_loss: accept when the sample-weighted global loss difference is negative. Set loss_weight to the rank's local sample count.

Both branches are averaged through PyTorch's native model-averaging utilities. Loss-only closure evaluations run under torch.no_grad(), so the closure must guard backward work with torch.is_grad_enabled() to avoid extra DDP gradient synchronization. See Distributed training for the execution sequence, policy semantics, and integration requirements.

• AADL.reset_acceleration_history(optimizer): clear history after any external parameter-changing operation.
• AADL.remove_acceleration(optimizer): restore the original optimizer step and remove AADL state.
• AADL.accept_candidate(...): low-level scalar policy reducer for advanced integrations that already provide comparable candidate and baseline losses.

Pull requests are welcome. For major changes, please open an issue first to discuss what you would like to change.

BSD-3-Clause

M. Lupo Pasini, V. Reshniak, and M. K. Stoyanov. AADL: Anderson Accelerated Deep Learning. Computer Software. https://github.com/ORNL/AADL.git. 06 Sep. 2021. Web. doi:10.11578/dc.20210723.1. Copyright ID#: 81927550

M. Lupo Pasini, J. Yin, V. Reshniak and M. K. Stoyanov, "Anderson Acceleration for Distributed Training of Deep Learning Models," SoutheastCon 2022, 2022, pp. 289-295, doi: 10.1109/SoutheastCon48659.2022.9763953.