CASMO: Confident Adaptive Selective Momentum Optimizer

The Optimizer That Knows What To Learn.

Abstract

CASMO is a novel optimization algorithm designed to address the limitations of uniform gradient application in deep learning. Unlike standard optimizers (e.g., AdamW, SGD) that treat all gradients as equally valid, CASMO introduces Adaptive Gradient-Aware Regularization (AGAR). This mechanism dynamically quantifies the "signal-to-noise" ratio of gradients over time, allowing the optimizer to selectively dampen updates that correspond to noise, outliers, or conflicting information while accelerating learning on generalizable patterns.

This approach yields state-of-the-art stability in challenging regimes such as label noise, long-tail distributions, differential privacy, and continual learning.

💡 Key Innovation: The AGAR Mechanism

The core of CASMO is the Adaptive Gradient Alignment Ratio (AGAR), a metric that measures the consistency of gradient direction over a sliding window:

$$ \text{AGAR} = \frac{||\mathbb{E}[g]||^2}{||\mathbb{E}[g]||^2 + \text{Var}[g]} $$

• High AGAR ($\approx 1$): Indicates consistent, reliable signal (e.g., generalizable features). CASMO applies full updates.
• Low AGAR ($\approx 0$): Indicates high variance or conflict (e.g., label noise, privacy noise, catastrophic forgetting). CASMO automatically dampens the learning rate.

This allows CASMO to act as an automatic filter during training, improving robustness without manual hyperparameter tuning or data filtering.

⚡ Quick Start

Installation

Via PyPI (Recommended):

pip install casmo-optimizer

From Source (For Development):

git clone https://github.com/abderahmane-ai/CASMO.git
cd CASMO
pip install -e .

Usage

CASMO is a drop-in replacement for torch.optim.AdamW.

from casmo import CASMO

# Initialize with standard parameters
optimizer = CASMO(
    model.parameters(), 
    lr=1e-3, 
    weight_decay=0.01,
    granularity='group'  # Recommended for efficiency
)

# Standard PyTorch training loop
for batch in dataloader:
    optimizer.zero_grad()
    loss = model(batch)
    loss.backward()
    optimizer.step()

🧪 Benchmarks & Reproducibility

We provide a comprehensive suite of benchmarks demonstrating CASMO's superiority in specific failure modes of standard optimizers. Each benchmark contains full reproduction scripts and detailed analysis.

See benchmarks/README.md for detailed results and reproduction instructions.

🛠️ Configuration

Citation

If you use CASMO in your research, please cite:

@software{casmo2025,
  title={CASMO: Confident Adaptive Selective Momentum Optimizer},
  author={Ainouche, Abderahmane},
  year={2025},
  url={https://github.com/abderahmane-ai/CASMO}
}

Made with ❤️ for the ML research community.