🎓 A1Learner
🔍

PyTorch A Custom Gradient, Checked

A straight-through estimator is a custom function with a forward that rounds and a backward that copies. The copy pretends rounding never happened, so a gradient can flow through a step that has none.

Plain rounding blocks the gradient

Rounding snaps 0.3 down and 0.7 up. Its gradient prints zero for both inputs, so no learning signal survives.

import torch

x = torch.tensor([0.3, 0.7], requires_grad=True)
y = torch.round(x)
print("Rounded:", y.tolist())
y.sum().backward()
print("Gradient from round:", x.grad.tolist())

The estimator passes the gradient through

RoundSTE rounds in forward, exactly like plain rounding. Its backward returns the incoming gradient unchanged, so both inputs receive one.

class RoundSTE(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        return torch.round(x)

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output


x = torch.tensor([0.3, 0.7], requires_grad=True)
y = RoundSTE.apply(x)
print("Rounded:", y.tolist())
y.sum().backward()
print("Gradient through the estimator:", x.grad.tolist())

One step of learning signal

The rounded prediction misses its target of one. The gradient for w still prints minus two, so an optimizer can move w.

w = torch.tensor(0.2, requires_grad=True)
score = w * torch.tensor(1.0)
pred = RoundSTE.apply(score)
loss = (pred - torch.tensor(1.0)) ** 2
loss.backward()
print("Prediction:", pred.item())
print("Loss:", loss.item())
print("Gradient for w:", w.grad.item())

Example

import torch


class RoundSTE(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        return torch.round(x)

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output


torch.manual_seed(0)
X = torch.tensor([-2.0, -1.0, 1.0, 2.0])
targets = torch.tensor([0.0, 0.0, 1.0, 1.0])

w = torch.tensor(0.1, requires_grad=True)
b = torch.tensor(0.0, requires_grad=True)
optimizer = torch.optim.Adam([w, b], lr=0.1)
loss_fn = torch.nn.MSELoss()

for step in range(60):
    raw = w * X + b
    pred = RoundSTE.apply(raw)
    loss = loss_fn(pred, targets)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if step in (0, 20, 40, 59):
        print(f"Step {step}: loss {loss.item():.4f}")

with torch.no_grad():
    final_pred = RoundSTE.apply(w * X + b)
print("Learned w:", round(w.item(), 3))
print("Learned b:", round(b.item(), 3))
print("Final predictions:", [int(v) for v in final_pred.tolist()])
print("Targets:", [int(v) for v in targets.tolist()])

Key points

  • Rounding gives a gradient of zero almost everywhere, so plain rounding cannot train
  • The estimator keeps rounding in forward and replaces only the backward rule
  • In backward, the estimator returns the incoming gradient unchanged
  • Threshold model: a model whose answer flips when a score crosses a fixed line
  • In short, the forward rounds honestly and the backward lies usefully

In this example, RoundSTE rounds each raw score in forward and copies the gradient in backward. The data holds four inputs with targets zero, zero, one, and one. The model is one weight and one bias, so each raw score is w times x plus b. Adam trains the pair for 60 steps at rate 0.1. The loss prints 0.5000 at step 0 and 0.2500 at step 20. By step 40 the loss prints 0.0000, and it stays there. The learned w prints 0.239 and the learned b prints 0.471. The final predictions print as 0, 0, 1, 1 and match the targets. Plain rounding could never take these steps, because its gradient is zero.

Note

The training result is the check here, because the estimator gradient is invented by design

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why can the estimator train when plain rounding cannot?