Batching and broadcasting

A model could process one example at a time, but in practice you feed it a batch: many examples stacked as rows of one matrix. A batch turns many matrix-vector products into a single matrix-matrix product, and that keeps a GPU busy. Broadcasting is the rule that lets you add a bias vector to every row of the batch without copying it by hand.

The big idea

One example is a matrix-vector product. Stack B examples as rows of X, which has shape (B, d_in). Then X · W is one matrix multiply, and the same weights are reused for every row. The weights are loaded once and used many times, which is why GPUs handle batches so well.

Walking through it

Here is a batched linear layer. It runs as written.

import numpy as np

X = np.array([              # batch of 4 examples, 3 features each
    [2.0, 1.0, 0.5],
    [0.0, 1.0, 1.0],
    [1.0, 0.0, 0.0],
    [3.0, 2.0, 1.0],
])                          # shape (4, 3)

W = np.array([              # shape (3, 2): in_features × out_features
    [0.5, 1.0],
    [-1.0, 0.0],
    [2.0, 1.0],
])

b = np.array([0.1, -0.2])   # shape (2,)

Y = X @ W + b               # (4, 2): the bias is broadcast across rows
print(Y.shape)
print(Y[0])

Output:

(4, 2)
[1.1 2.3]

The first row checks out by hand. [2, 1, 0.5] · [0.5, −1, 2] = 1, plus the bias 0.1 gives 1.1. The second column gives 2.5, minus 0.2, gives 2.3.

Broadcasting

Broadcasting combines arrays of different shapes. It lines up the shapes from the right. For each pair of dimensions, they must be equal, or one must be 1, or one must be missing.

(4, 2) + (2,)    → (4, 2)   bias added to every row
(4, 1) * (1, 3)  → (4, 3)   outer-product style
(4, 3) + (4,)    → error    (4,) lines up with the last dim 3, and 3 ≠ 4
(4, 3) + (4, 1)  → (4, 3)   each row gets its own value

Broadcasting does not copy memory. Attention masks rely on it too: a (1, 1, T, T) mask applies to every example and every head in a batch.

Common mistakes

• Dropping the batch dimension. A vector of shape (3,) and a batch of shape (1, 3) behave differently in some operations. Check the shape you actually pass.
• Trusting a broadcast that works. If it runs without error, it may still be doing the wrong thing. Assert shapes in critical code.
• Mixing weight conventions. X · W expects W as (in, out). nn.Linear stores (out, in) and computes X · Wᵀ + b.

Check yourself

1. What is the shape of X · W when X is (32, 784) and W is (784, 256)? Answer: (32, 256). The inner 784s match, and the batch size stays.
2. Does (4, 3) + (4,) work? Answer: No. The last dimensions are 3 and 4, and neither is 1.

Key takeaways

• A batch of B examples is one matrix multiply with shape (B, d_in).
• Weights are reused across the batch, which is why GPUs are fast.
• Broadcasting adds biases and masks without copying data.
• Broadcasting can fail silently when the shapes happen to line up.

Further learning

Next, count the cost of these multiplies: FLOP counting for matmul.