popriskmin gives you PRM, a small modification of AdamW that trains on population risk instead of raw empirical risk. It keeps AdamW's usual one forward pass and one backward pass per step, then preconditions each parameter update with the population-risk mask from Litman & Guo (2026), A Theory of Generalization in Deep Learning.

PRM tracks one extra tensor per parameter tensor: an exponential moving average of centered minibatch gradient variance. For each scalar parameter $k$, it asks whether the batch-mean gradient is larger than the leave-one-out noise estimate:

$$ \mu_k^2 > \alpha \sigma_k^2 $$

Parameters that pass get the Adam update. Parameters that fail get shrunk or zeroed, depending on the mask. On the fresh-batch boundary, $\alpha = 1$ and the streaming variance estimates the leave-one-out penalty $\Sigma_B / (b - 1)$.

uv pip install -e .
uv sync --extra test

PRM requires torch>=2.0.

from popriskmin import PRM

optimizer = PRM(
    model.parameters(),
    lr=3e-4,
    weight_decay=0.01,
    softness=1.0,
    batch_size=32,
)

for batch in loader:
    optimizer.zero_grad()
    loss = loss_fn(model(batch.x), batch.y)
    loss.backward()
    optimizer.step()

Useful options:

PRM(model.parameters(), mask="snr")         # default, smooth SNR mask
PRM(model.parameters(), mask="soft")        # strict Algorithm 1 cutoff
PRM(model.parameters(), mask="hard")        # 0/1 theorem mask
PRM(model.parameters(), reduction="per_tensor")

Use reduction="per_tensor" when each scalar parameter is noisy but the whole parameter tensor has a clear signal. This is often the more useful setting for large generative models, diffusion, and CFM-style training.

softness is $\lambda_p$. Larger values make the mask more conservative. A reasonable first sweep is 0.3, 1, 3, 10.

stats = optimizer.get_mask_stats()
print(stats)

Example output:

{
    "mean_q": 0.62,
    "active_fraction": 0.71,
    "min_q": 0.00,
    "max_q": 0.99,
    "parameter_count": 1_245_184,
    "noise_scale": 12.4,
    "signal_sq": 18.7,
    "snr": 1.51,
}

mean_q tells you how open the mask is on average. If active_fraction falls near zero, almost every parameter is below the leave-one-out boundary. Try a lower softness, a higher learning rate, or reduction="per_tensor" before assuming the method is broken.

PRM is meant for settings where plain empirical-risk training fits structured noise or memorizes before it generalizes. The paper reports improvements in:

With more experiments in the appendix. However, PRM is not magic. If AdamW already reaches the solution without overfitting, PRM may not help.

Run the smoke test:

uv run python examples/synthetic_regression.py

The script trains AdamW and several PRM variants on a noisy regression problem. You should see the optimizer learning the underlying signal instead of the corrupted training labels.

popriskmin/
|-- popriskmin/
|   |-- __init__.py
|   |-- mask.py
|   `-- optimizer.py
|-- examples/
|   `-- synthetic_regression.py
|-- tests/
|   `-- test_optimizer.py
`-- pyproject.toml

@misc{litman2026theory,
  title         = {A Theory of Generalization in Deep Learning},
  author        = {Litman, Elon and Guo, Gabe},
  year          = {2026},
  eprint        = {2605.01172},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  doi           = {10.48550/arXiv.2605.01172}
}