Building a Layer

A layer is a pile of neurons that all read the same input and write their results side by side. That single change turns a loop over neurons into one matrix multiply, which is exactly why GPUs are useful at all. A layer you can write in three lines is a layer you can debug.

The idea

Stack out neurons, each with its own weights and bias, and the layer is:

• W has one row per neuron, so W alone is the whole layer of weights.
• X W^T contracts the in axis for every sample in the batch at once, giving (B, out).
• The transpose exists only because PyTorch stores W as (out, in). X @ W would try to contract the batch axis and fail.
• The bias (out,) broadcasts across the batch.

The parameter count is out × in + out. A Linear(768, 3072) feed-forward projection holds 768·3072 + 3072 = 2,362,368 numbers.

Initialization sets the training dynamics. If weight entries have standard deviation s, pre-activations of a zero-mean input grow with in·s². PyTorch draws Linear weights from U(−1/√in, 1/√in), which is std ≈ 1/√(3·in) — for in = 768 that is 0.0208. Kaiming initialization roughly doubles that for ReLU layers to compensate for the clipped negative half.

Worked example

With X of shape (2, 3) and a Linear(3, 2) layer:

For the first sample and first neuron: 0.1·1.0 + (−0.2)·0.0 + 0.3·(−1.0) = −0.2, plus 0.05 gives −0.15, and ReLU clips it to 0. The full output is [[0.0, 0.9], [0.15, 0.05]].

In code

import torch

X = torch.tensor([[1.0, 0.0, -1.0], [0.5, 0.5, 0.5]])
W = torch.tensor([[0.1, -0.2, 0.3], [0.4, 0.5, -0.6]])
b = torch.tensor([0.05, -0.1])

manual = torch.relu(X @ W.T + b)
layer = torch.nn.Linear(3, 2)
with torch.no_grad():
    layer.weight.copy_(W)
    layer.bias.copy_(b)

print(manual)                       # tensor([[0.0000, 0.9000], [0.1500, 0.0500]])
print(torch.allclose(manual, torch.relu(layer(X))))          # True
print(sum(p.numel() for p in torch.nn.Linear(768, 3072).parameters()))   # 2362368
print(torch.nn.Linear(768, 3072).weight.std().item())        # ~0.0208 = 1/sqrt(3*768)

Check yourself

1. X is (8, 768) and the layer is Linear(768, 3072). What is the output shape, and how many parameters does the layer hold?
2. Why write X @ W.T + b rather than X @ W + b?
3. You initialize W with standard deviation 10. What happens to activations and to the first gradients, and why?

Key takeaways

• A layer is a stack of neurons, so its forward pass is one matmul plus a broadcast bias.
• The (out, in) weight layout is why the matmul needs a transpose.
• Initialization scale — around 1/√in — is what keeps activation variance stable across depth.