🎓 A1Learner
🔍

PyTorch Scaling Features

Feature scaling rewrites each feature column so the columns share a similar size range. Standardising is one scaling method, built from each column's mean and standard deviation.

One step, two gradient sizes

The pain comes first. One backward pass on raw columns gives the two weights wildly different gradients.

torch.manual_seed(1)
small = torch.randn(50, 1)
big = torch.randn(50, 1) * 50
x = torch.cat([small, big], dim=1)
y = 2.0 * small + 0.08 * big + 0.1 * torch.randn(50, 1)
torch.manual_seed(0)
model = nn.Linear(2, 1)
loss = nn.MSELoss()(model(x), y)
loss.backward()
grads = model.weight.grad.flatten()
print("Gradient for the small feature:", round(grads[0].item(), 2))
print("Gradient for the big feature:", round(grads[1].item(), 1))

Standardise one column

Subtract the mean, then divide by the std. The column lands on mean 0 and std 1.

column = torch.tensor([10.0, 20.0, 30.0, 40.0])
scaled = (column - column.mean()) / column.std()
print("Scaled:", [round(v, 2) for v in scaled.tolist()])
print("Mean after:", round(scaled.mean().item(), 4))
print("Std after:", round(scaled.std().item(), 4))

Example

import torch
from torch import nn


def fit(x, y, lr, steps=80):
    torch.manual_seed(0)
    model = nn.Linear(2, 1)
    loss_fn = nn.MSELoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
    trace = []
    for step in range(steps):
        pred = model(x)
        loss = loss_fn(pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if step % 20 == 0 or step == steps - 1:
            trace.append(round(loss.item(), 4))
    return trace


torch.manual_seed(1)
small = torch.randn(50, 1)
big = torch.randn(50, 1) * 50
x = torch.cat([small, big], dim=1)
y = 2.0 * small + 0.08 * big + 0.1 * torch.randn(50, 1)

mean = x.mean(dim=0)
std = x.std(dim=0)
scaled = (x - mean) / std

print("Column std before scaling:", [round(v, 1) for v in x.std(dim=0).tolist()])
print("Raw scale losses:", fit(x, y, lr=0.0001))
print("Column std after scaling:", [round(v, 2) for v in scaled.std(dim=0).tolist()])
print("Scaled losses:", fit(scaled, y, lr=0.1))

Key points

  • Mean: the average value of a column
  • Std: short for standard deviation, a measure of a column's spread
  • Standardised value: (value - mean) / std, worked out cell by cell
  • A standardised column has mean 0 and std 1
  • Use dim=0 in mean() and std() to work down each column
  • Wildly different scales force one tiny learning rate onto every weight
  • Scaling gives every column the same spread, so one rate fits all

In this example, seed 1 makes two feature columns of 50 rows. One column spreads near 1 and the other near 50. The targets mix both columns. The fit function trains a Linear model with SGD. Seed 0 is set inside fit before each model is built. Both fits therefore start from the same weights. The raw fit must use rate 0.0001, since a bigger rate explodes on the big column. Its loss drops once, then crawls from 4.2492 to 4.1503. The mean and std of each column then standardise the data. Both column stds print as 1.0 after scaling. The scaled fit uses rate 0.1 and reaches 0.0102. Same data, same model, and the scaled columns train cleanly.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does standardising a column do to its mean and std?