Models From First Principles 07: PACS — Building an Optimizer From Gradient Statistics
PACS — Building an Optimizer From Gradient Statistics
So far in Models From First Principles, every post has asked some version of the same question:
What should the model compute?
MR.Q gave us a scalar quality estimate.
EBT split one shared representation into Q, V and Policy.
SICQL made those heads explicit, replaceable components.
HRM introduced repeated computation over fast and slow latent states.
Tiny compressed iterative refinement into one recursive latent state.
Then we opened Tiny and found residual blocks, attention and sparse autoencoders inside it.
Now we move underneath all of them.
Because regardless of how clever the architecture is, training eventually reaches the same point:
loss
↓
backward()
↓
parameter gradients
↓
optimizer
↓
new parameters
A model can have recurrence, attention, uncertainty heads, sparse concepts and hierarchical state.
But if its parameters do not move in a useful direction, none of that matters.
This post builds a small custom optimizer called PACS from first principles.
The important idea is not the name.
The important idea is that an optimizer is also just a model of sorts: it keeps state, observes gradients, transforms those gradients and decides how much each parameter should move.
Our version will maintain two statistics per parameter:
current gradient
│
├───────────────┐
↓ ↓
mean gradient mean squared gradient
│ │
│ adaptive scale
│ │
└──────┬────────┘
↓
preconditioned update
↓
parameter
By the end of the post, you should be able to read a custom PyTorch optimizer without treating optimizer.step() as magic.
1. Start with the simplest possible optimizer
Suppose our model has one scalar parameter:
import torch
w = torch.tensor([5.0], requires_grad=True)
And our loss is:
loss = (w - 2.0).pow(2).mean()
Then:
loss.backward()
print(w.grad)
The derivative is:
d/dw (w - 2)^2 = 2(w - 2)
At w = 5:
gradient = 6
The most basic gradient-descent update is:
w_new = w_old - learning_rate × gradient
In PyTorch:
with torch.no_grad():
w -= 0.1 * w.grad
That is an optimizer.
Everything else in this post is a transformation applied between:
raw gradient
and:
parameter update
2. SGD is already a policy for interpreting gradients
The phrase “gradient descent” can make the update sound inevitable.
It is not.
The gradient tells us the local direction of steepest increase of the loss.
The optimizer chooses what to do with that information.
Plain SGD says:
update_t = g_t
θ_{t+1} = θ_t - η g_t
where:
θis the parameter;ηis the learning rate;g_tis the current gradient.
A minimal PyTorch optimizer could therefore look like this:
from torch.optim import Optimizer
class PlainSGD(Optimizer):
def __init__(self, params, lr=1e-3):
super().__init__(params, {"lr": lr})
@torch.no_grad()
def step(self):
for group in self.param_groups:
lr = group["lr"]
for p in group["params"]:
if p.grad is None:
continue
p.add_(p.grad, alpha=-lr)
The optimizer does not need to understand attention, recurrence or language.
It sees tensors and gradients.
That separation is one of PyTorch’s most useful abstractions.
3. Why use gradient history?
Real training gradients are noisy.
Mini-batches differ.
Examples conflict.
Multi-task heads pull shared parameters in different directions.
A single gradient can therefore be a poor estimate of the direction we want to follow over many batches.
One simple response is to maintain an exponential moving average:
m_t = β m_{t-1} + (1 - β) g_t
where β might be 0.9.
The optimizer is now remembering recent gradient direction.
In code:
m.mul_(beta).add_(grad, alpha=1 - beta)
This is a very small piece of code.
Conceptually, however, it changes the optimizer from:
react to the current batch
into:
combine the current batch with recent training history
4. A toy noisy-gradient example
Imagine the gradients for one parameter are:
+1.2
+0.8
+1.1
-0.4
+0.9
+1.0
The fourth gradient points in the opposite direction.
With no memory, that batch immediately reverses the update direction.
With a moving average, its effect is moderated by the previous history.
Let’s make that visible:
gradients = torch.tensor([1.2, 0.8, 1.1, -0.4, 0.9, 1.0])
beta = 0.9
m = torch.tensor(0.0)
for g in gradients:
m = beta * m + (1 - beta) * g
print(f"grad={g.item():+.3f} avg={m.item():+.3f}")
The moving average does not tell us which gradient is “correct.”
It encodes an assumption:
recent gradient direction contains useful signal that should persist across batches.
That assumption must be tested.
5. Direction is not the only problem
Suppose our model has two parameters whose gradients typically look like this:
parameter A: 0.0008
parameter B: 14.0
Using one global learning rate means those parameters can experience very different effective update scales.
One solution is to track the squared gradient:
v_t = ρ v_{t-1} + (1 - ρ) g_t²
Then normalize the update by:
sqrt(v_t) + ε
This produces a diagonal preconditioner.
Large historical squared gradients increase the denominator.
Small historical squared gradients produce a smaller denominator.
The optimizer therefore adapts its step independently for each parameter element.
6. Build the preconditioner
The update is:
v_t = ρ v_{t-1} + (1 - ρ) g_t²
In PyTorch:
precond.mul_(rho).addcmul_(
grad,
grad,
value=1 - rho,
)
Then:
denom = precond.sqrt().add(eps)
And a normalized gradient would be:
scaled_grad = grad / denom
Notice what happened.
We started with:
parameter ← raw gradient
and now have:
parameter
↑
scaled gradient
↑
historical squared gradients
The optimizer now has internal state.
7. Combine direction memory and adaptive scaling
PACS combines the two ideas:
g_t
│
├───────────────┐
↓ ↓
m_t v_t
│ │
│ sqrt(v_t)+ε
│ │
└──────┬────────┘
↓
m_t / sqrt(v_t)
↓
× η
↓
parameter update
Mathematically:
m_t = β m_{t-1} + (1 - β) g_t
v_t = ρ v_{t-1} + (1 - ρ) g_t²
u_t = m_t / (sqrt(v_t) + ε)
θ_{t+1} = θ_t - η u_t
That is already enough to implement the optimizer.
8. Build PACS as a real PyTorch optimizer
from __future__ import annotations
from typing import Optional
import torch
from torch.optim import Optimizer
class PACSOptimizer(Optimizer):
def __init__(
self,
params,
lr: float = 1e-4,
beta: float = 0.9,
eps: float = 1e-8,
weight_decay: float = 0.0,
preconditioner_decay: float = 0.999,
):
defaults = dict(
lr=lr,
beta=beta,
eps=eps,
weight_decay=weight_decay,
preconditioner_decay=preconditioner_decay,
)
super().__init__(params, defaults)
@torch.no_grad()
def step(self, closure: Optional[callable] = None):
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
for group in self.param_groups:
lr = group["lr"]
beta = group["beta"]
eps = group["eps"]
weight_decay = group["weight_decay"]
preconditioner_decay = group["preconditioner_decay"]
for p in group["params"]:
if p.grad is None:
continue
grad = p.grad
if weight_decay != 0:
grad = grad.add(p, alpha=weight_decay)
state = self.state[p]
if len(state) == 0:
state["step"] = 0
state["grad_avg"] = torch.zeros_like(p)
state["precond"] = torch.zeros_like(p)
grad_avg = state["grad_avg"]
precond = state["precond"]
state["step"] += 1
grad_avg.mul_(beta).add_(grad, alpha=1 - beta)
precond.mul_(preconditioner_decay).addcmul_(
grad,
grad,
value=1 - preconditioner_decay,
)
denom = precond.sqrt().add(eps)
update = grad_avg / denom
p.add_(update, alpha=-lr)
return loss
That is the core PACS optimizer.
No hidden training framework is required.
9. What state does the optimizer actually store?
For every parameter tensor, PACS stores:
step
mean gradient
mean squared gradient
If a parameter is shaped:
[256, 256]
then both optimizer-state tensors are also:
[256, 256]
This matters for memory.
A model with P parameters does not necessarily require only P parameter values during training.
PACS adds approximately two state values per trainable parameter:
parameters P
mean gradient P
preconditioner P
before counting gradients, activations and other training state.
Optimizer choice therefore affects memory as well as convergence.
10. Inspect optimizer memory directly
def optimizer_state_elements(optimizer):
total = 0
for state in optimizer.state.values():
for value in state.values():
if torch.is_tensor(value):
total += value.numel()
return total
After one step:
print(optimizer_state_elements(optimizer))
For PACS, expect roughly:
2 × trainable parameter count
in tensor-valued optimizer state.
That is an observable property, not an implementation detail to ignore.
11. Weight decay: where we apply it matters
The implementation above uses:
grad = grad.add(p, alpha=weight_decay)
which is equivalent to adding an L2-style parameter penalty into the gradient before the moving averages are updated.
This is not the same thing as decoupled weight decay.
A decoupled update would instead do something like:
p.mul_(1 - lr * weight_decay)
separately from the gradient transformation.
Those two choices can behave differently when the optimizer adaptively rescales gradients.
So when comparing optimizers, do not write simply:
weight_decay=0.01
and assume the meaning is identical.
The implementation defines the semantics.
12. PACS is related to existing optimizer ideas — not magically outside them
The PACS update contains familiar ingredients:
- first-moment smoothing;
- second-moment smoothing;
- elementwise adaptive scaling;
- optional weight decay.
Those ingredients also appear in established optimizers.
That does not make every such optimizer identical.
Important differences can include:
- exact recurrence equations;
- bias correction;
- where epsilon is placed;
- coupled vs decoupled weight decay;
- momentum interpretation;
- initialization;
- parameter groups;
- learning-rate schedules.
This version of PACS does not perform first- or second-moment bias correction.
That has consequences early in training.
13. Why initialization biases the first steps
Both moving averages start at zero:
m_0 = 0
v_0 = 0
At the first step:
m_1 = (1 - β) g_1
v_1 = (1 - ρ) g_1²
Those are biased toward zero because the history before step one was initialized to zero rather than sampled from the true gradient process.
Some optimizers explicitly correct this.
Our current PACS implementation does not.
That is neither automatically good nor automatically bad.
It is a property to measure.
14. Inspect the effective update
Looking only at gradient norm is not enough once the optimizer transforms gradients.
What we really want to inspect is:
raw gradient
moving-average gradient
preconditioner
final update
A diagnostic helper:
def pacs_parameter_stats(optimizer):
rows = []
for group_index, group in enumerate(optimizer.param_groups):
for p in group["params"]:
if p.grad is None:
continue
state = optimizer.state.get(p, {})
grad_avg = state.get("grad_avg")
precond = state.get("precond")
row = {
"group": group_index,
"shape": tuple(p.shape),
"grad_norm": float(p.grad.norm()),
}
if grad_avg is not None:
row["grad_avg_norm"] = float(grad_avg.norm())
if precond is not None:
row["precond_mean"] = float(precond.mean())
row["precond_max"] = float(precond.max())
denom = precond.sqrt().add(group["eps"])
update = grad_avg / denom
row["update_norm"] = float(update.norm())
rows.append(row)
return rows
This lets us ask the right question:
What update did the optimizer actually construct from the gradient?
15. Verify that parameters actually move
The debugging rule from the PyTorch series still applies.
Never stop at:
gradient exists
Prove the parameter changed.
def snapshot_parameters(model):
return {
name: p.detach().clone()
for name, p in model.named_parameters()
}
def parameter_change(before, model):
out = {}
for name, p in model.named_parameters():
out[name] = float(
(p.detach() - before[name]).norm()
)
return out
Then:
before = snapshot_parameters(model)
loss.backward()
optimizer.step()
print(parameter_change(before, model))
This is especially useful with custom optimizers because a bug in state initialization, parameter registration or update logic can silently leave some parameters unchanged.
16. Build a tiny benchmark problem
Use a problem where we know the model should learn.
import torch
import torch.nn as nn
torch.manual_seed(7)
x = torch.randn(2048, 8)
true_w = torch.randn(8, 1)
y = x @ true_w + 0.05 * torch.randn(2048, 1)
model = nn.Sequential(
nn.Linear(8, 64),
nn.GELU(),
nn.Linear(64, 1),
)
Training helper:
def train(model, optimizer, steps=500):
losses = []
for _ in range(steps):
pred = model(x)
loss = ((pred - y) ** 2).mean()
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
losses.append(float(loss.detach()))
return losses
Now PACS can be tested rather than discussed abstractly.
17. Compare against baselines
A custom optimizer should not be evaluated in isolation.
At minimum compare against:
torch.optim.SGD
torch.optim.RMSprop
torch.optim.Adam
torch.optim.AdamW
But there is a trap.
Using the same numerical learning rate for every optimizer is not necessarily a fair comparison.
Different update rules can have very different sensible learning-rate ranges.
A more meaningful experiment sweeps learning rates for every optimizer.
For example:
1e-5
3e-5
1e-4
3e-4
1e-3
3e-3
1e-2
Then compare each optimizer near its own best stable region.
18. What should we measure?
Do not measure only final training loss.
Useful metrics include:
best validation loss
steps to threshold
wall-clock time to threshold
stability across seeds
peak memory
optimizer-state memory
update norm
parameter norm
frequency of non-finite values
sensitivity to learning rate
An optimizer can reach a good loss while being:
- slower;
- more memory hungry;
- less stable;
- more sensitive to hyperparameters.
Those are real trade-offs.
19. Learning curves matter
Suppose two optimizers produce:
Optimizer A:
step 100: 0.20
step 200: 0.10
step 500: 0.08
Optimizer B:
step 100: 0.50
step 200: 0.20
step 500: 0.07
Which is better?
There is no universal answer.
If you stop at 200 steps, A wins.
If final convergence dominates, B might win.
If B costs twice as much memory, the answer can change again.
Optimization is an engineering decision, not a leaderboard number detached from constraints.
20. Compare steps to a target loss
def first_below(losses, threshold):
for i, loss in enumerate(losses):
if loss <= threshold:
return i + 1
return None
Then compare:
first_below(losses, 0.1)
This often communicates optimization behavior more clearly than one final number.
21. Gradient averaging as a variance-reduction hypothesis
PACS describes its first-moment accumulator as variance reduction.
That claim can be tested.
Collect gradients for a fixed parameter over many mini-batches:
g_1, g_2, ..., g_n
Then compare the variance of:
raw gradients
against:
moving-average gradients
For example:
raw = torch.stack(raw_grad_vectors)
avg = torch.stack(avg_grad_vectors)
print(raw.var(dim=0).mean())
print(avg.var(dim=0).mean())
If the second quantity is lower, smoothing reduced variability.
But lower variance alone does not prove better optimization.
The smoothed direction may also lag behind a changing objective.
Again:
mechanism ≠ benefit
22. Measure gradient/update alignment
We can ask how strongly the transformed update agrees with the raw gradient.
import torch.nn.functional as F
alignment = F.cosine_similarity(
grad.flatten(),
update.flatten(),
dim=0,
)
Near 1 means the optimizer mostly preserves the raw direction.
Smaller values mean historical state and preconditioning are changing the update substantially.
That is useful telemetry.
23. Preconditioning changes relative coordinate scale
Imagine:
grad = [0.01, 10.0]
Plain SGD preserves that 1000× magnitude difference.
An adaptive preconditioner can reduce it dramatically because the second coordinate also accumulates a much larger squared-gradient history.
This is one reason adaptive optimizers can work well when parameter dimensions experience very different gradient scales.
But it is also why they can behave differently from SGD even when the gradient directions appear similar.
24. Inspect effective learning rates
For each coordinate, PACS roughly applies:
η_eff = η / (sqrt(v_t) + ε)
We can inspect that directly:
effective_lr = group["lr"] / (
precond.sqrt() + group["eps"]
)
Then log:
print(
effective_lr.min().item(),
effective_lr.mean().item(),
effective_lr.max().item(),
)
A single configured learning rate therefore becomes a distribution of effective coordinate-level scales.
25. The epsilon is not decorative
The denominator is:
sqrt(v_t) + ε
If v_t is tiny, epsilon prevents division by zero.
But epsilon also affects scaling when the second moment is very small.
Values such as:
1e-8
1e-6
1e-4
can therefore alter behavior in low-gradient regimes.
For most ordinary training, you should not tune epsilon before the major hyperparameters.
But when implementing an optimizer from scratch, it is important to understand why it exists.
26. The two decay rates control different memories
PACS has:
beta
for the gradient average, and:
preconditioner_decay
for squared-gradient history.
They should not be conceptually collapsed.
A typical setup might be:
beta = 0.9
preconditioner_decay = 0.999
The first statistic adapts relatively quickly.
The second remembers scale over a longer horizon.
This creates two temporal filters inside the optimizer.
Interestingly, we have returned to an idea that appeared in HRM:
multiple timescales
But now the timescales belong to optimization statistics, not latent model state.
27. Sweep the two timescales independently
A useful experiment:
beta ∈ {0.0, 0.5, 0.9, 0.99}
rho ∈ {0.9, 0.99, 0.999, 0.9999}
Measure:
- convergence speed;
- validation performance;
- update volatility;
- run-to-run stability.
This tells us whether both memories are actually contributing.
28. Ablation: remove gradient averaging
Set:
beta = 0
Then:
m_t = g_t
The optimizer becomes an adaptive second-moment method without first-moment smoothing.
If performance is unchanged, the gradient-average state may not be buying much for that task.
29. Ablation: remove preconditioning
Replace:
update = grad_avg / denom
with:
update = grad_avg
Now we isolate the contribution of historical gradient smoothing.
This is closer to momentum-style optimization.
Again, test rather than assume.
30. Ablation: raw SGD
Remove both memories:
update = grad
Now our hierarchy is:
SGD
↓
smoothed gradient
↓
smoothed + preconditioned gradient
This is the optimizer equivalent of the model decompositions we have used throughout the series.
31. PACS as composition
The complete optimizer is not one indivisible idea.
It is:
PACS
├── gradient observation
├── first-moment accumulator
├── second-moment accumulator
├── elementwise preconditioner
├── weight-decay rule
└── parameter update
Each component can be tested independently.
That is the central theme of Models From First Principles again.
32. Multi-task models make optimizer behavior especially interesting
Consider the models from earlier posts.
A shared encoder may receive gradients from:
Q loss
V loss
policy loss
uncertainty loss
reconstruction loss
consistency loss
The resulting gradient is an aggregate of several objectives.
PACS then applies temporal smoothing and coordinate-wise scaling to that aggregate.
This means optimizer state can interact with multi-task conflict.
A direction that appears briefly in one task might be suppressed by history.
A persistent direction might accumulate influence.
Whether that helps is empirical.
33. Inspect per-loss gradients before the optimizer sees them
For a shared parameter p, compute task-specific gradients:
g_q = torch.autograd.grad(
q_loss,
p,
retain_graph=True,
)[0]
g_v = torch.autograd.grad(
v_loss,
p,
retain_graph=True,
)[0]
Then compare:
F.cosine_similarity(
g_q.flatten(),
g_v.flatten(),
dim=0,
)
This separates two questions:
- are the objectives conflicting?
- how does the optimizer transform their combined gradient?
Do not blame the optimizer for conflict created by the objective.
34. Optimizer state is part of the checkpoint
If you save only:
torch.save(model.state_dict(), "model.pt")
then resume training with a fresh PACS optimizer, you have discarded:
gradient averages
preconditioner history
step counters
The model parameters are restored.
The optimization process is not.
A training checkpoint should therefore include:
torch.save(
{
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"step": global_step,
},
"checkpoint.pt",
)
35. Verify optimizer checkpoint restoration
After loading:
ckpt = torch.load("checkpoint.pt")
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["optimizer"])
inspect several state entries.
The goal is to prove that:
preconditioner_before_save
≈
preconditioner_after_load
and similarly for the first moment.
Checkpoint tests should include optimizer state, not just model outputs.
36. Replacing model parameters can invalidate optimizer state
This is the same stale-optimizer problem we discussed in the SICQL post.
Suppose:
optimizer = PACSOptimizer(model.parameters())
Then later:
model.score_head = nn.Linear(256, 1)
The optimizer still owns the old parameter objects.
The new head will not automatically appear in its parameter groups.
After structural replacement, rebuild or explicitly update the optimizer.
37. Parameter groups are first-class architecture
Different model components may need different optimization settings.
For example:
optimizer = PACSOptimizer(
[
{
"params": model.encoder.parameters(),
"lr": 1e-4,
},
{
"params": model.score_head.parameters(),
"lr": 5e-4,
},
]
)
This is useful when:
- one component is pretrained;
- one head is newly initialized;
- some components are fragile;
- different modules learn at different speeds.
But every extra hyperparameter increases experimental degrees of freedom.
Use groups when evidence justifies them.
38. A numerical-safety audit
Custom optimizers should explicitly check for non-finite gradients.
if not torch.isfinite(grad).all():
raise RuntimeError("non-finite gradient")
You may also inspect:
if not torch.isfinite(update).all():
raise RuntimeError("non-finite PACS update")
This catches cases where:
- gradients explode;
- state becomes corrupted;
- denominator scaling produces invalid values;
- upstream loss becomes NaN.
39. Do not silently mutate .data unless you understand why
Older optimizer code often uses:
p.data.add_(...)
Inside an optimizer’s @torch.no_grad() step, you can usually write:
p.add_(...)
instead.
That makes the intent clearer and avoids unnecessary .data usage.
The essential requirement is that optimizer updates should not themselves become part of the autograd graph.
40. A safer PACS implementation
Here is a slightly hardened version:
class PACS(Optimizer):
def __init__(
self,
params,
lr=1e-4,
beta=0.9,
rho=0.999,
eps=1e-8,
weight_decay=0.0,
):
if lr <= 0:
raise ValueError("lr must be positive")
if not 0 <= beta < 1:
raise ValueError("beta must be in [0,1)")
if not 0 <= rho < 1:
raise ValueError("rho must be in [0,1)")
if eps <= 0:
raise ValueError("eps must be positive")
defaults = dict(
lr=lr,
beta=beta,
rho=rho,
eps=eps,
weight_decay=weight_decay,
)
super().__init__(params, defaults)
@torch.no_grad()
def step(self, closure=None):
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
for group in self.param_groups:
for p in group["params"]:
if p.grad is None:
continue
g = p.grad
if not torch.isfinite(g).all():
raise RuntimeError("PACS received non-finite gradient")
if group["weight_decay"] != 0:
g = g.add(
p,
alpha=group["weight_decay"],
)
state = self.state[p]
if not state:
state["step"] = 0
state["m"] = torch.zeros_like(p)
state["v"] = torch.zeros_like(p)
m = state["m"]
v = state["v"]
state["step"] += 1
m.mul_(group["beta"]).add_(
g,
alpha=1 - group["beta"],
)
v.mul_(group["rho"]).addcmul_(
g,
g,
value=1 - group["rho"],
)
denom = v.sqrt().add_(group["eps"])
update = m / denom
if not torch.isfinite(update).all():
raise RuntimeError("PACS produced non-finite update")
p.add_(update, alpha=-group["lr"])
return loss
Nothing in this code is inaccessible.
The optimizer is just stateful tensor arithmetic.
41. Test the optimizer on one parameter first
Before training a deep model, test PACS on a scalar quadratic.
w = torch.nn.Parameter(torch.tensor([10.0]))
optimizer = PACS([w], lr=0.05)
for step in range(200):
loss = (w - 3.0).pow(2).mean()
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
print(w.item())
You should end near:
3.0
If the optimizer cannot solve this, do not debug it inside an HRM.
42. Test tensor shapes next
Use parameters of different shapes:
p1 = nn.Parameter(torch.randn(10))
p2 = nn.Parameter(torch.randn(4, 8))
p3 = nn.Parameter(torch.randn(2, 3, 5))
After a step, verify:
assert optimizer.state[p1]["m"].shape == p1.shape
assert optimizer.state[p2]["v"].shape == p2.shape
assert optimizer.state[p3]["m"].shape == p3.shape
Optimizer state must mirror parameter shape exactly.
43. Test missing gradients
Some parameters may be unused in a particular forward pass.
Then:
p.grad is None
should simply skip the parameter.
That is different from:
p.grad == 0
A zero gradient is a real tensor-valued result.
A missing gradient means autograd did not produce one for that parameter in the current backward pass.
44. Test repeated zero_grad
The correct training pattern remains:
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
If you forget to clear gradients, PyTorch accumulates them.
Then PACS receives the accumulated gradient rather than the current mini-batch gradient.
The optimizer cannot distinguish whether that accumulation was intentional.
45. Gradient accumulation changes optimizer frequency
Suppose you accumulate gradients across four mini-batches before calling step().
Then optimizer state updates once per four batches rather than once per batch.
This changes the temporal meaning of:
beta
rho
because one optimizer step now represents more training examples.
Hyperparameters expressed per optimizer step do not automatically preserve their behavior when accumulation strategy changes.
46. Mixed precision adds another layer
With automatic mixed precision, gradients may be scaled before unscaling.
A custom optimizer should receive unscaled gradients.
Typical CUDA training uses a scaler workflow where unscaling occurs before the optimizer update.
When debugging PACS under mixed precision, inspect gradients after unscaling, otherwise the apparent norms can be misleading.
The safest sequence is:
forward under autocast
↓
scaled backward
↓
unscale gradients
↓
optional clipping / diagnostics
↓
optimizer step
↓
scaler update
47. Gradient clipping and PACS
There are two conceptually different places you could clip:
raw gradient
or:
transformed update
Most ordinary PyTorch workflows clip gradients before optimizer.step().
That means PACS sees the clipped gradient and updates both moving statistics from it.
If you instead clip the final preconditioned update, you are defining a different optimizer.
Document which one you mean.
48. Time the optimizer step separately
For large models, optimizer cost itself can matter.
Measure:
import time
start = time.perf_counter()
optimizer.step()
elapsed = time.perf_counter() - start
On CUDA, remember that operations are asynchronous, so use proper synchronization or CUDA events when you need accurate device timing.
Then compare optimizer-step overhead separately from forward/backward cost.
49. Benchmark state initialization separately
The first call to step() creates optimizer state lazily.
That first step can therefore be more expensive than later steps.
Do not benchmark only step one and call it steady-state optimizer performance.
Warm up first.
50. Measure optimizer sensitivity, not just its best run
Suppose PACS reaches a good score only at:
lr = 0.00031
and fails badly around it.
Another optimizer achieves nearly the same score over:
0.0001 → 0.003
The second optimizer may be more useful in practice even if its best run is microscopically worse.
Plot performance across hyperparameter ranges.
Robustness is part of optimizer quality.
51. Run multiple seeds
Optimization results can be noisy.
At minimum compare several seeds:
1
2
3
4
5
Report:
mean
standard deviation
min/max
failure count
One lucky run is weak evidence for a new optimizer.
52. Separate architecture gains from optimizer gains
This is particularly important for our series.
Suppose Tiny with PACS beats HRM with AdamW.
That does not tell us whether Tiny is better than HRM.
It mixes two interventions:
architecture
optimizer
A better matrix is:
AdamW PACS
MR.Q ✓ ✓
SICQL ✓ ✓
HRM ✓ ✓
Tiny ✓ ✓
Now optimizer and architecture effects can be separated.
53. Fair optimizer experiments are expensive
For each optimizer you ideally tune:
- learning rate;
- weight decay;
- momentum/first moment;
- second-moment decay;
- scheduler interaction;
- clipping threshold.
This can become a large search.
That is why optimizer claims should be scoped carefully.
A result such as:
PACS was better under this exact configuration and budget
is much stronger than:
PACS is a better optimizer.
54. A useful experiment report
For every optimizer experiment, record:
model architecture
parameter count
dataset version
embedding version
objective
batch size
gradient accumulation
learning-rate schedule
optimizer hyperparameters
seed
hardware
number of steps
wall-clock time
validation metric
peak memory
Without this context, reproducing optimization results becomes difficult.
55. What PACS actually guarantees
The implementation guarantees that it:
- keeps an exponential moving average of gradients;
- keeps an exponential moving average of squared gradients;
- scales the smoothed gradient by the square root of the second moment;
- applies the result to parameters using a configured learning rate.
It does not guarantee:
- faster convergence;
- better generalization;
- lower gradient variance in every regime;
- better handling of multi-task conflict;
- more stable recursion;
- superior calibration.
Those are experimental questions.
56. This is the optimizer version of our evidence rule
Throughout this series we have repeated:
head name ≠ proven semantic
The equivalent rule here is:
optimizer mechanism ≠ proven training benefit
Moving averages exist.
Preconditioning exists.
Whether those mechanisms improve your actual model must be measured.
57. The deeper model-inside-the-model idea
We started the series with:
model
↓
smaller models
↓
blocks
↓
layers
↓
tensor operations
PACS extends that stack downward:
model
↓
loss
↓
gradients
↓
optimizer state
↓
transformed update
↓
parameter mutation
The training algorithm itself can be decomposed.
58. Final compact implementation
Here is the complete optimizer again without commentary:
import torch
from torch.optim import Optimizer
class PACS(Optimizer):
def __init__(
self,
params,
lr=1e-4,
beta=0.9,
rho=0.999,
eps=1e-8,
weight_decay=0.0,
):
defaults = dict(
lr=lr,
beta=beta,
rho=rho,
eps=eps,
weight_decay=weight_decay,
)
super().__init__(params, defaults)
@torch.no_grad()
def step(self, closure=None):
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
for group in self.param_groups:
for p in group["params"]:
if p.grad is None:
continue
g = p.grad
if group["weight_decay"] != 0:
g = g.add(
p,
alpha=group["weight_decay"],
)
state = self.state[p]
if not state:
state["step"] = 0
state["m"] = torch.zeros_like(p)
state["v"] = torch.zeros_like(p)
state["step"] += 1
m = state["m"]
v = state["v"]
m.mul_(group["beta"]).add_(
g,
alpha=1 - group["beta"],
)
v.mul_(group["rho"]).addcmul_(
g,
g,
value=1 - group["rho"],
)
denom = v.sqrt().add_(group["eps"])
update = m / denom
p.add_(update, alpha=-group["lr"])
return loss
That is it.
The optimizer is no longer a black box.
59. Where the series goes next
We have now travelled from model outputs all the way down to parameter mutation:
MR.Q
↓
EBT
↓
SICQL
↓
HRM
↓
Tiny
↓
Tiny internals
↓
PACS
The next step is to zoom back out.
We now have enough machinery to compare the architectures as a family rather than as isolated posts.
The next article will ask:
What does each architecture actually add, what does it cost, and under what evidence would we choose one over another?
That is where the series becomes a model-selection framework rather than a collection of implementations.
