The Fully Connected Layer as an Affine Map

A dense layer exists to apply one affine transformation to a batch of vectors. The transposed weight layout, the bias vector that gets broadcast along the batch axis, the number of stored parameters, and the shape of every gradient that flows back through it are consequences of that single identity plus a few layout conventions. Getting the shapes and the reduction conventions wrong does not usually crash anything, which is why they are worth understanding separately from the arithmetic of training.

The map and the storage convention

For an input matrix $X$ of shape $(N, n_{\text{in}})$, a layer with weight $W$ of shape $(n_{\text{out}}, n_{\text{in}})$ and bias $b$ of shape $(n_{\text{out}},)$ computes

$$Y = X W^{\top} + b$$

Written per element, $y_{ni} = \sum_{j=1}^{n_{\text{in}}} x_{nj} W_{ij} + b_i$. Row $i$ of $W$ is the coefficient vector of output unit $i$, and the transpose exists because storage is indexed by output unit first while the matrix product needs the contraction over $n_{\text{in}}$ to be an inner product. Nothing in that expression couples row $n$ to row $n’$: the layer sees a batch as $N$ independent vectors that happen to travel together through the hardware.

The bias is the one part of the affine map that is not a matrix, and broadcasting it along the leading axis fixes what it can express. The same $b$ is added to every sample, so a layer cannot use its bias to record per-sample offsets, and shifting every input by a constant vector $c$ shifts every output by $Wc$. Adding a bias to a layer also makes the composition sensitive to affine rescalings of its input in a way a purely linear layer is not, which matters for the collapse described below.

Parameter count, arithmetic, and batch behaviour

The stored quantities are $n_{\text{in}} n_{\text{out}} + n_{\text{out}}$ numbers, all of which are trained. A layer mapping 784 inputs to 1000 units holds $784000 + 1000 = 785000$ parameters, about 3.0 MiB in 32-bit floats. Because each output element needs $n_{\text{in}}$ multiply-accumulate operations, a forward pass over a batch of $N$ costs $N n_{\text{in}} n_{\text{out}}$ multiply-accumulates, or roughly $2 N n_{\text{in}} n_{\text{out}}$ floating-point operations once the addition is counted.

The backward pass produces three quantities: the gradient with respect to the weights, which has the shape of $W$ and equals $\mathrm{d}Y^{\top} X$; the gradient with respect to the input, which has the shape of $X$ and equals $\mathrm{d}Y,W^{\top}$; and the gradient with respect to the bias, which is $\mathrm{d}Y$ summed over the batch axis, because broadcasting is a sum run backwards. Two matrix products instead of one make the backward pass cost about twice the forward pass in multiply-accumulates, a ratio that holds for this layer regardless of its size. The input gradient is only needed when another layer sits below; the weight gradient is always needed, and the weight gradient needs $X$, so an input tensor that would otherwise be discardable is retained until the backward pass reaches that layer. For a network holding $L$ dense layers, activations dominate memory at large batch sizes; at small batch sizes the weights dominate, and the arithmetic intensity of the layer is limited by the weight matrix being streamed once per batch. A batch of one reads $n_{\text{in}} n_{\text{out}}$ floats to perform $n_{\text{in}} n_{\text{out}}$ multiply-accumulates, while a batch of a thousand reuses each weight a thousand times.

That reuse is also why the layer is executed as a general matrix-matrix product rather than as a loop, and the transposed layout is what the level-3 BLAS interface expects: its general routine computes $C \leftarrow \alpha A B + \beta C$ with an option to transpose either operand, so a weight matrix stored output-first maps onto it without a copy. Storing the weights with the input dimension first would force either an explicit transpose on every call, which costs a second pass over the matrix, or a strided access pattern for the multiply. The convention is set by the interface of the library performing the arithmetic rather than by any property of the mathematics. It also means the cost model of a dense layer inherits the sensitivity of matrix multiplication to shape: the same number of multiply-accumulates costs more when both operands are long and thin, because there is less opportunity to keep an operand resident in fast memory while it is reused across the other dimension. Batching improves that reuse, and the batch size that can be used is bounded by the memory the activations occupy rather than by the arithmetic.

Two affine layers are one affine layer

Substituting one layer into the next, for a single input vector $x$, gives

$$W_2(W_1 x + b_1) + b_2 = (W_2 W_1), x + (W_2 b_1 + b_2)$$

The composition is again an affine map, with an effective matrix that is the product of the two, so its rank cannot exceed $\min(n_{\text{in}}, n_{\text{hidden}}, n_{\text{out}})$. Stacking more affine maps only produces a longer factorization of the same single matrix, and depth buys nothing beyond the rank limit imposed by the narrowest layer. A hidden layer narrower than the input is an information bottleneck of that rank, and no later affine layer can restore what it discarded.

Removing the bias does not change the argument; it makes the layer linear rather than affine, and every stacked pair still collapses. The consequence is that the expressive power of a deep stack comes entirely from whatever is inserted between the affine maps, which is the subject of pointwise nonlinearity. An affine layer on its own can only rotate, scale, shear, and translate.

import numpy as np

def dense_forward(X, W, b):          # X: (N, n_in), W: (n_out, n_in), b: (n_out,)
    return X @ W.T + b

def dense_backward(X, W, dY):        # dY: (N, n_out)
    dW = dY.T @ X                    # contraction over the batch axis
    db = dY.sum(axis=0)              # the broadcast add in reverse
    dX = dY @ W                      # only needed if a layer sits below
    return dX, dW, db

X = np.random.randn(4, 3)
W = np.random.randn(5, 3)
b = np.zeros(5)
Y = dense_forward(X, W, b)
dX, dW, db = dense_backward(X, W, np.ones_like(Y))
assert dW.shape == W.shape and dX.shape == X.shape and db.shape == b.shape

Reduction conventions and other ways the gradient goes wrong

The bias gradient and the weight gradient have different natural reduction conventions, and mixing them is the most common silent defect in a hand-written layer. With a loss that is a mean over a batch of $N$, the correct bias gradient is $\frac{1}{N}\sum_n \mathrm{d}y_{ni}$, while the weight gradient is $\frac{1}{N}\sum_n \mathrm{d}y_{ni} x_{nj}$. Implementations that apply the $1/N$ factor in the loss function and none in the layer are consistent; implementations that apply it in one place and not the other are also consistent as long as the convention is uniform. Applying it to the weights and not the biases scales the bias step by $N$, which is invisible at batch size 1 and destabilizing at batch size 512. A gradient check against finite differences will show the discrepancy, as described in gradient checking.

A second defect is shape-correct transposition. Computing $\mathrm{d}W = X^{\top}\mathrm{d}Y$ yields a matrix whose shape matches neither $W$ nor $X$ unless the layer is square, so it usually fails loudly. Computing $\mathrm{d}X = \mathrm{d}Y,W$ instead of $\mathrm{d}Y,W^{\top}$ succeeds whenever $n_{\text{in}} = n_{\text{out}}$ and produces a consistent but wrong gradient.

When the same weight matrix is used in more than one place, as in a tied embedding, the two gradient contributions must be summed rather than overwritten. The general rule belongs to the reverse sweep of the computational graph, where a value with $k$ consumers receives $k$ incoming cotangents that add, which is the mechanism covered in reverse-mode differentiation.

Degenerate inputs and dead configurations

Several configurations produce a layer that is well-formed and untrainable. If a row of $W$ and its bias are both exactly zero for some unit, then $y_i = 0$ for every input, and the unit’s weight gradient is $\mathrm{d}y_i$ times the input, which is also zero, so nothing moves and the unit never recovers. Weight decay on a row that has already reached zero applies a zero gradient there too. If $n_{\text{in}} = 0$, the sum over inputs is empty and the layer reduces to the bias; an empty batch produces empty gradients and, in many frameworks, an average over zero samples. None of these are errors at the level of the affine map; they are errors at the level of the use of the parameters, and their symptom is a training loss that stops moving rather than an exception.