adam-atan2-pytorch

What it is and what it does

Adam-atan2-pytorch implements a modified Adam optimizer that replaces the epsilon term with atan2 to achieve numerical stability without requiring manual epsilon tuning. The core change is algorithmic—using atan2 in the update rule removes the need for epsilon altogether while maintaining scale invariance across different parameter ranges. This is based on research from DeepMind proposing the atan2 modification.

The package also includes features for improving plasticity in continual learning scenarios, where models must learn new tasks without catastrophic forgetting of old ones. It integrates with PyTorch's standard optimizer interface, so you instantiate it with model parameters and a learning rate, then call step() and zero_grad() in your training loop like any other optimizer. The implementation depends on einops for tensor operations and torch itself.

Use it for:

• Replace standard Adam in deep learning models where epsilon tuning is problematic or parameter scales vary widely.
• Continual learning pipelines where maintaining plasticity across task sequences is critical.
• Training scenarios where numerical stability without manual epsilon adjustment is desired.
• Research comparing optimizer variants with scale-invariant update rules.
• Large-scale training where removing epsilon hyperparameter reduces tuning overhead.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Provides an Adam optimizer variant using atan2 for numerical stability and scale invariance, with optional plasticity features for continual learning.

Yes, if you use PyTorch and want to experiment with a theoretically motivated Adam variant. Low install friction, permissive license, no known vulnerabilities, and active maintenance make it a low-risk addition. Best suited for researchers exploring optimizer variants or teams dealing with scale-invariance issues; standard Adam remains the safer default for most production training.

Install

adam-atan2-pytorch on PyPI

pip install adam-atan2-pytorch

uv add adam-atan2-pytorch

poetry add adam-atan2-pytorch

Installing adam-atan2-pytorch

Before you install

Low install friction with only two runtime dependencies (einops and torch). Active maintenance status with a recent release 28 days ago.

License in practice

MIT License permits unrestricted use, modification, and distribution with only attribution and liability disclaimer required—no restrictions on commercial or private use.

Quickstart

pip install adam-atan2-pytorch

import torch
from adam_atan2_pytorch import AdamAtan2

model = torch.nn.Linear(10, 1)
opt = AdamAtan2(model.parameters(), lr=1e-4)
loss = model(torch.randn(10))
loss.backward()
opt.step()
opt.zero_grad()

Requires PyPI installation of torch and einops; Python 3.9 or later.

Verify before relying

• Whether plasticity features are enabled by default or require explicit configuration.
• Performance comparison with standard Adam on typical training tasks.
• Specific continual learning scenarios where plasticity features provide measurable benefit.

Package facts

Evidence: adam_atan2_pytorch-0.4.0-py3-none-any.whl

Keywords: adam, artificial intelligence, deep learning, optimizers

Tags

More Artificial Intelligence packages

Further reading