stac-optimizer

Korean README | Optimizer docs | Korean docs | Benchmark JSON

STAC means "SignSGD Trunk, AdamW Cap". It keeps the sign trunk state-free, uses AdamW only on the final trainable-module tail, and is tuned to reduce optimizer-state VRAM without giving up tail stability.

Flow

flowchart LR
    A["Trainable modules<br/>registration order"]
    B["Resolve AdamW cap<br/>`last_n_modules` or<br/>default `last_n_ratio=12.5%`"]

    subgraph S["State-free sign trunk"]
        C["Earlier modules"]
        D["Decoupled weight decay on weight tensors<br/>bias + 1-D params skip decay by default<br/>parameter -= lr * sign(grad)<br/>no momentum, no sign-side state"]
    end

    subgraph T["AdamW cap"]
        E["Final tail modules"]
        F["Standard AdamW on the tail<br/>bias + 1-D params skip decay by default<br/>exp_avg + exp_avg_sq"]
    end

    A --> B
    B --> C
    B --> E
    C --> D
    E --> F

    classDef neutral fill:#f8fafc,stroke:#475569,color:#0f172a,stroke-width:1px;
    classDef sign fill:#d7f0e8,stroke:#0f766e,color:#134e4a,stroke-width:1.5px;
    classDef adam fill:#dbeafe,stroke:#2563eb,color:#1d4ed8,stroke-width:1.5px;

    class A,B neutral;
    class C,D sign;
    class E,F adam;

Install

python -m pip install stac-optimizer

For local development and benchmark generation:

python -m pip install -e ".[dev]"

Quickstart

import torch
from torch import nn

from stac_optimizer import STAC


model = nn.Sequential(
    nn.Linear(128, 64),
    nn.ReLU(),
    nn.Linear(64, 32),
    nn.ReLU(),
    nn.Linear(32, 10),
)

optimizer = STAC(
    model,
    lr=1e-3,
    last_n_ratio=0.125,
    weight_decay=1e-2,
    error_if_nonfinite=True,
)

loss = torch.nn.functional.mse_loss(
    model(torch.randn(8, 128)),
    torch.randn(8, 10),
)
loss.backward()
optimizer.step()
optimizer.zero_grad(set_to_none=True)

last_n_ratio counts only modules that directly own trainable parameters. Pure containers such as nn.Sequential are skipped unless they own parameters themselves. Use last_n_modules when you want an explicit cap size instead. Bias tensors and 1-D parameters such as LayerNorm scales skip decoupled weight decay by default in both sections.

CUDA Research Snapshot

The repository benchmark is CUDA-only and uses held-out validation splits, 5 paired seeds, seeded teachers, seeded student initialization, fixed batch schedules per seed, deep residual models, a transformer-like sequence task with embeddings and LayerNorm, BF16 autocast when supported, epoch-by-epoch validation loss curves, and a first-step optimizer-memory probe. The AdamW baseline uses the same bias/1-D no-decay grouping so the comparison does not hinge on a different weight-decay policy.

Snapshot from 2026-03-19 on torch 2.10.0+cu126 and NVIDIA GeForce RTX 3070:

Repository takeaway: the default preset cuts optimizer state from 98.227 MB to 8.133 MB, the full-decay variant keeps the same memory profile while helping the norm-heavy and sequence tasks a bit, and the wider cap spends more AdamW state to improve regression. Those are repository-local measurements, not universal guarantees.

Verify

python -m pytest -q
python examples/research_benchmark.py --device cuda
rm -rf build dist
python -m build
python -m twine check dist/*