Metadata-Version: 2.4
Name: temperon
Version: 0.1.0
Summary: Temperon: pay for sharpness-aware minimization only where it pays you back
Project-URL: Homepage, https://github.com/MStamatis/temperon
Project-URL: Repository, https://github.com/MStamatis/temperon
Project-URL: Issues, https://github.com/MStamatis/temperon/issues
Author-email: Stamatis Mastromichalakis <stamatis@tmnetworks.gr>
License-Expression: MIT
License-File: LICENSE
Keywords: muon,optimizer,pytorch,sam,sharpness-aware-minimization,training-efficiency
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Requires-Dist: torch>=2.0
Description-Content-Type: text/markdown

# Temperon

**Pay for sharpness-aware minimization only where it pays you back.**

*Tempering* is the controlled heat treatment that follows quenching and sets a
metal's final toughness. This package governs the training equivalent: what
happens during the last anneal.

[SAM](https://arxiv.org/abs/2010.01412) improves generalization but doubles the
cost of every training step, which is why most people skip it. Temperon runs SAM
for a contiguous tail aligned to your learning-rate decay instead of the whole
run. In our experiments that reaches full-time-SAM quality for roughly a third
less wall-clock — and spreading the same SAM budget uniformly across training
does *worse* than either extreme.

> Status: **pre-release** (`0.1.0.dev0`), published alongside a paper in
> preparation. The API may still change. Single-GPU, bf16/fp32 tested;
> DDP and gradient accumulation are not supported yet.

```bash
pip install temperon
```

## Usage

```python
import torch
from temperon import Temperon, wsd

base = torch.optim.AdamW(model.parameters(), lr=2e-5)
total = len(loader) * epochs

opt = Temperon(base, total_steps=total, tail_frac=0.3, rho=0.05)
sched = torch.optim.lr_scheduler.LambdaLR(base, wsd(total, 250, decay_frac=0.3))

for batch in loader:
    def closure():
        opt.zero_grad()
        loss = loss_fn(model(batch.x), batch.y)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        return loss

    loss = opt.step(closure)   # 1 pass in the cheap phase, 2 in the tail
    sched.step()
```

The closure follows the `torch.optim.LBFGS` convention: zero the gradients,
compute the loss, call `backward()`, return the loss. Temperon calls it once per
step during the cheap phase and twice during the tail. Nothing else in your
training code changes.

### The one thing that matters

`tail_frac` must be **aligned to your learning-rate decay**, so the tail owns a
complete anneal. This is not a detail: in our ablations a SAM tail bolted onto
the middle of an ongoing cycle gained nothing at all (flat in tail length,
~1pp below full SAM), while a tail owning a fresh anneal matched full-time SAM.
Pass the same fraction to `wsd(..., decay_frac=f)` and `Temperon(tail_frac=f)`
and the alignment is exact.

### Choosing `tail_frac`

The switch point is a measured quality/cost dial, not a magic number. On
CIFAR-100 (5 seeds, same 100-epoch budget, SGD explorer → SAM+Muon tail):

| `tail_frac` | final accuracy | total cost (calibrated) |
|---|---|---|
| 0.57 | 0.8295 ± 0.0034 — ties full-time SAM+Muon | 5038s |
| 0.20 | 0.8249 ± 0.0022 — −0.42pp vs full SAM+Muon (p=0.015), level with the published full SAM+SGD recipe (+0.17pp, p=0.35) | 2862s, **−43%** |

A longer tail buys the expensive method's full ceiling at its full cost; a
shorter one gives up a measured slice of that ceiling for a much cheaper run.
The default 0.3 is the value that matched full-time SAM quality on GPT-2
pretraining (−29% wall-clock) and GLUE fine-tuning (−36%). Below 0.2 we have
no measurements; as `tail_frac → 1` you are simply paying for full-time SAM.

### Cheap explorer, expensive refiner

You can also switch optimizers at the same boundary — a cheap explorer for most
of training, then a stronger optimizer owning the SAM tail:

```python
from temperon import cosine_tail

explorer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
refiner  = Muon(model.parameters(), lr=0.01, weight_decay=0.2)  # any Muon impl.

opt = Temperon(explorer, total_steps=total, tail_frac=0.57,
             tail_optimizer=refiner, transfer="momentum")
sched = torch.optim.lr_scheduler.LambdaLR(explorer, cosine_tail(total, 0.57))
```

`transfer="momentum"` carries `momentum_buffer` across the switch.

## When *not* to use this

Measured limits, stated because they define where the method applies:

- **A cheap optimizer wins below its own ceiling.** Temperon never accelerates an
  accuracy target that plain SGD or Adam can already reach; it accelerates
  targets that only expensive methods reach at all. If your target is modest,
  use the cheap optimizer.
- **Saturated tasks gain nothing.** Where the frontier is reached early, the
  late anneal is not what is limiting you.
- **Language-model pretraining is out of scope.** In near-single-epoch
  pretraining there is little overfitting for SAM to prevent, and we measured
  no useful gain; under heavy data repetition SAM was actively worse. SAM's
  established language gains are in *fine-tuning*, which is where this package
  is aimed for LMs.

## Citing

See [CITATION.cff](https://github.com/MStamatis/temperon/blob/main/CITATION.cff).
The research code and the full experimental record — including the negative
controls this claim rests on — live in the same repository.

## License

MIT.
