Skip to content

nactttch/TTA

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

4 Commits
 
 
 
 
 
 
 
 
 
 

Repository files navigation

TwoTimescaleAgree (TTA) — A Geometric Momentum Optimizer

PyTorch License: MIT

TwoTimescaleAgree (TTA) is a production‑grade PyTorch optimizer that uses agreement‑gated two‑timescale momentum to break the AdamW(lr × c) symmetry. It was the only non‑degenerate mechanism identified among the methods explored in the geometry study26.3 % of its update coordinates flip sign relative to Adam, something a positive gain (scalar or per‑coordinate) can never do.

⚠️ HONEST STATUS
On the controlled tasks we could run, TTA’s ~10 % edge over AdamW was statistically indistinguishable from simply using a ~1.5 × larger learning rate.
TTA is delivered as a geometric, minimal‑overhead candidate to validate with a fair (best‑lr‑per‑optimizer, multi‑seed) benchmark — NOT as an established AdamW‑beater.

Current Evidence

Evidence collected so far:

  • Geometry study (sign‑flip analysis):

    • 26.3 % of update coordinates flip sign relative to Adam.
    • Gain‑based variants (AgreeAdam, SNRGate) produced 0.0 % sign flips.
  • Small‑scale image classification benchmark (20 epochs, 5 seeds, test error %):

    Optimizer LR=0.003 (best)
    AdamW 22.32% ± 0.35%
    TTA 24.94% ± 0.31%
    Relative improvement: ~+11.7 % (lower error).
  • Validation‑loss experiment (single seed, longer run):

    Optimizer Validation loss Wall time
    AdamW 2.8847 17.7 min
    TTA 2.7071 18.1 min

These experiments are encouraging but insufficient to establish superiority over AdamW. Large‑scale, multi‑seed benchmarks with independent hyper‑parameter tuning are still necessary.

Why TwoTimescaleAgree?

  • Geometric non‑degeneracy – Instead of scaling a single direction (like AgreeAdam/SNRGate, which measured exactly 0.0 % sign flips), TTA switches between two different directional estimates (fast vs. slow momentum) depending on whether they agree in sign. This is what changes the optimization geometry.
  • Minimal overhead – Only 3 state buffers (fast EMA, slow EMA, second moment), one more than AdamW, with no reductions, no matrix multiplications.
  • Decoupled weight decay – Weight decay is applied exactly as in AdamW.

Algorithm

For each parameter coordinate:

  1. Maintain fast and slow exponential moving averages of the gradient:
    m_fast ← β_fast · m_fast + (1 - β_fast) · g
    m_slow ← β_slow · m_slow + (1 - β_slow) · g   (β_slow > β_fast)
    
  2. Maintain the second moment:
    v ← β₂ · v + (1 - β₂) · g²
    
  3. Bias‑correct all buffers and compute the denominator:
    ĥ_fast = m_fast / (1 - β_fastᵗ),   ĥ_slow = m_slow / (1 - β_slowᵗ)
    denom = √(v / (1 - β₂ᵗ)) + ε
    
  4. Agreement gate selects the update direction:
    direction = ĥ_fast   if sign(ĥ_fast) == sign(ĥ_slow)
                ĥ_slow   otherwise
    
  5. Apply decoupled weight decay and the parameter update:
    p ← p - lr · weight_decay · p
    p ← p - lr · direction / denom
    

This element‑wise switching (not scaling) is what makes the method fundamentally different from Adam‑style gain‑based approaches.

Installation

Just copy two_timescale_agree.py into your project.
Dependencies: PyTorch ≥ 1.10.

pip install torch   # if not already installed

Usage

from two_timescale_agree import TwoTimescaleAgree

optimizer = TwoTimescaleAgree(
    model.parameters(),
    lr=1e-3,                       # tune per task
    betas=(0.9, 0.99, 0.999),      # (β_fast, β_slow, β₂)
    eps=1e-8,
    weight_decay=0.01,
)

Full API matches torch.optim.Optimizer (including maximize, closures, etc.). See the docstring in the code for details.

How You Can Help – Call for Tests & Benchmarks

I do not own a powerful GPU cluster and cannot run large‑scale, multi‑seed, best‑lr‑per‑optimizer benchmarks. If you have the compute, please help!

What would be valuable

  1. Fair comparison

    • Tune the learning rate (and possibly β_fast, β_slow) independently for each optimizer on a validation set.
    • Run multiple seeds (≥ 5) and report mean ± std.
    • Compare against AdamW, Lion, SGD‑M, Adan, Sophia, etc.
  2. Diverse workloads

    • NLP: fine‑tuning BERT/GPT‑2, training small transformers from scratch.
    • Vision: ResNet, ViT on ImageNet‑1k or CIFAR‑10/100.
    • RL: PPO / DQN on standard Gym environments.
    • GANs / Diffusion models: where optimiser choice is critical.
  3. Hyper‑parameter sensitivity

    • β_fast, β_slow grid search (typical ranges: 0.8‑0.99).
    • Interaction with weight decay, batch size, scheduler.
  4. Long‑run stability

    • Train for hundreds of epochs and check if TTA diverges or plateaus differently.

How to contribute

  • Run the optimizer on your own tasks and open an issue with the results (a simple text table is enough!).
  • If you find interesting behaviours (better convergence, worse stability, etc.), please share the logs and hyper‑parameters.
  • Pull requests with improvements to the code, documentation, or test scripts are very welcome.

Contact:

Citation

If you use TTA in your research, please cite this repository. A BibTeX entry will be added once a paper is available; for now, you can cite:

@misc{TwoTimescaleAgree,
  author = {Ahmed},
  title = {TwoTimescaleAgree: Agreement-gated Two-Timescale Momentum Optimizer},
  year = {2025},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/nactttch/TTA}}
}

License

This project is licensed under the MIT License – see the LICENSE file for details.


Happy optimising! Let’s find out together whether TTA can reliably outperform AdamW.

About

PyTorch implementation of TwoTimescaleAgree (TTA): an agreement-gated two-timescale momentum optimizer that breaks AdamW symmetry with minimal overhead. Decoupled weight decay.

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages