DouglasOrr
commented
2 months ago Hi @norikazu99, thanks for your comments & for sharing your question.
The mse loss doesn't seem to scale_fwd whereas cross entropy loss does. On the other hand mse loss scales bwd for both input and target by a fixed value, while cross entropy doesn't. Can you please explain the reasons behind these decisions.
I'll have a go at annotating the code. An important thing to note is that for losses, we assume constraint=None, which means we're free to choose separate scaling factors in the forward and backward passes.
# Cross entropy loss
input = scale_bwd(input, vocab_size / (vocab_size - 1) ** 0.5) # from [unit scaling] Table 5
input = scale_fwd(input, mult) # multiplier to "sharpen" the softmax; no need to scale grad since constraint=None
loss = F.cross_entropy(...)
if reduction == "mean":
return scale_fwd(loss, 1 / batch_size) # just scale_fwd, thanks to separate scaling of fwd/bwd
assert reduction == "sum"
return loss
# MSE loss
grad_scale = 8**-0.5 # Assume `input` and `target` are independent ~N(0, 1) & grad(input)=2*(input-target)
input = scale_bwd(input, grad_scale)
target = scale_bwd(target, grad_scale) # typically target.requires_grad is False, but just in case
loss = F.mse_loss(input, target, size_average, reduce, reduction="sum")
if reduction == "mean":
return scale_fwd(loss, 1 / input.nelement()) # just scale_fwd, thanks to separate scaling of fwd/bwd
assert reduction == "sum"
return loss[unit scaling] https://arxiv.org/abs/2303.11257
So I think they're reasonably consistent. Softmax cross entropy has no scaling for labels which can't receive a gradient. MSE loss has no mult, as I think a multiplier on the inputs wouldn't make sense, and a multiplier on (inputs - targets) would end up as a scaling factor on both forward and backward passes, which we'd then cancel out to recover unit scale.
R.e. quantile loss, that looks broadly reasonable, yes, although I am unfamiliar with the particulars of this loss function. If independent Gaussian inputs are a reasonable assumption, it looks like the scale should possibly be $\times2$ not $\times2\sqrt{2}$, however:
quantiles = torch.tensor([0.5]) # multiple quantile scaling is broken
y = torch.randn(1000, 128, 1).requires_grad_()
y_hat = torch.randn(1000, 128, len(quantiles)).requires_grad_()
def loss_fn(y_hat, y):
loss = repeat(y, 'b s d -> b s (d n)', n=len(quantiles)) - y_hat
loss = torch.max((quantiles-1.0)*loss, quantiles*loss).mean()
return loss
def unit_mup_loss_fn(y_hat, y):
grad_scale = 2 # 8 ** -0.5
y_hat = scale_bwd(y_hat, grad_scale)
y = scale_bwd(y, grad_scale)
loss = repeat(y, 'b s d -> b s (d n)', n=len(quantiles)) - y_hat
loss = torch.max((quantiles-1.0)*loss, quantiles*loss).sum()
return scale_fwd(loss, 1 / y_hat.nelement())
y.grad = y_hat.grad = None
loss = loss_fn(y_hat, y)
loss.backward()
print(" unscaled", loss.item(), y.grad.std().item(), y_hat.grad.std().item())
y.grad = y_hat.grad = None
loss = unit_mup_loss_fn(y_hat, y)
loss.backward()
print("unit-scaled", loss.item(), y.grad.std().item(), y_hat.grad.std().item())This seems to be robust to different choices of a single quantile. But the scaling is not good when using multiple quantiles. This is because the repeat() causes a sum of gradients in the backward pass. This might require some thought about the distributions to fix. (Also, perhaps in the case of multiple quantiles, drawing y_hat independently is a bad idea?)
norikazu99
Hello and thanks a lot for sharing your great research and code.
The mse loss doesn't seem to scale_fwd whereas cross entropy loss does. On the other hand mse loss scales bwd for both input and target by a fixed value, while cross entropy doesn't. Can you please explain the reasons behind these decisions.
The following is my first attempt at implementing a unit_scaling version of quantile loss and would like to know if I missed something.