Tensors, Shapes, and Batching

Open the source of any transformer implementation and the comments that matter most are not about attention or normalisation, they are shape annotations: # (batch, seq, d_model). A tensor is just a multi-dimensional array of numbers, but the shape of that array is a contract every layer must honour. Get a shape wrong and the model either crashes with a dimension mismatch or, worse, silently broadcasts into nonsense. Reading transformer code fluently means reading shapes fluently.

What a tensor actually is

A tensor generalises scalars, vectors, and matrices to arbitrary dimension. A scalar is a 0-dimensional tensor, a vector is 1-dimensional, a matrix is 2-dimensional, and a transformer's activations are typically 3- or 4-dimensional. The rank (or number of dimensions) and the shape (the size along each dimension) together describe everything about a tensor's structure except the actual numbers inside it. A token embedding for one word is a vector of shape (d_model,), perhaps (4096,) for a large model. Stack the embeddings for every position in one sequence and you get a matrix of shape (seq_len, d_model). That is the object a single transformer layer actually operates on.

The batch dimension

Training and inference rarely process one sequence at a time; they process many in parallel to use hardware efficiently. This adds a leading batch dimension, so the tensor flowing through a transformer is usually shape (batch, seq_len, d_model). Every operation inside the model, matrix multiplication, normalisation, activation, is defined to act identically and independently across the batch dimension. This is what makes batching "free" in terms of code complexity: you write the layer for one example and the framework vectorises it across the batch axis, using the underlying hardware's parallelism (see broadcasting-and-vectorisation).

Attention adds one more wrinkle. Multi-head attention splits d_model into n_heads pieces of size d_head, so internally the tensor becomes (batch, n_heads, seq_len, d_head) before the attention computation and is reshaped back to (batch, seq_len, d_model) afterward. That reshape-split-merge pattern, not the attention maths itself, is where most beginner implementation bugs live.

Reading a shape trace

A useful habit: before reading any layer's forward pass, write down the shape of its input and predict the shape of its output. A linear layer with weight shape (d_in, d_out) maps (batch, seq_len, d_in) to (batch, seq_len, d_out), touching only the last dimension. A layer norm (see layernorm-residual-connections) preserves shape entirely, normalising within the last dimension. Attention is the one operation that mixes information across the seq_len axis; everything else in a transformer block operates position-wise, which is exactly why attention is the component that lets tokens exchange information at all (see attention-mechanism).

input:            (batch, seq_len, d_model)
after QKV proj:    (batch, seq_len, 3 * d_model)
split heads:       (batch, n_heads, seq_len, d_head)
after attention:   (batch, n_heads, seq_len, d_head)
merge heads:        (batch, seq_len, d_model)

When it falls down

• Silent broadcasting bugs. Two tensors with shapes that are technically compatible for broadcasting but semantically wrong (say, (seq_len,) accidentally broadcasting against (batch,) when both happen to be equal-sized) produce no error and wrong numbers. Shape mismatches that raise an exception are the easy case; shape coincidences that do not are the dangerous one.
• Variable-length sequences force padding decisions. Batching requires a rectangular tensor, so sequences of different lengths must be padded to a common length, and the padding must be masked out of both the loss and attention or it silently corrupts training.
• Memory scales with every dimension. A shape trace is also a memory estimate: doubling batch size or sequence length doubles activation memory for most tensors in the stack (see training-memory-footprint), which is why shape literacy is also cost literacy.
• Contiguity is invisible in a shape. Two tensors can report the identical shape while one is a memory-contiguous block and the other a strided view (e.g. after a transpose); some operations silently fail or need an explicit copy to run efficiently, and a shape alone will not tell you which case you are in.

Further reading

• PyTorch documentation, Tensor Views - how reshape, transpose, and view interact with memory layout.
• Vaswani et al., 2017, Attention Is All You Need, arXiv:1706.03762 - the original shapes: d_model, n_heads, d_k, d_v as defined for the first transformer.