Why nonlinearity — collapsing two linear layers

Background

The headline argument for activations: two stacked linear layers are just one linear layer. With no nonlinearity between them,

$$W_2(W_1\mathbf{x} + \mathbf{b}_1) + \mathbf{b}_2 = (W_2 W_1)\mathbf{x} + (W_2\mathbf{b}_1 + \mathbf{b}_2)$$

so the pair is equivalent to a single layer with weights $W = W_2 W_1$ and bias $\mathbf{b} = W_2\mathbf{b}_1 + \mathbf{b}_2$. Depth without nonlinearity buys nothing.

Problem statement

Implement collapse(W1, b1, W2, b2) returning the single equivalent layer (W, b).

Input

• W1 — shape (h, i), b1 — length h (first layer).
• W2 — shape (o, h), b2 — length o (second layer).

Output

Returns a tuple (W, b) where W is shape (o, i) and b is shape (o,), such that W @ x + b equals applying both linear layers in sequence.

Examples

Example 1

Input:  W1 = [[1, 2], [0, 1]], b1 = [1, 0], W2 = [[1, 0], [1, 1]], b2 = [0, 1]
Output: W = [[1, 2], [1, 3]],  b = [1, 2]

Explanation: $W_2W_1 = [[1,2],[1,3]]$ and $W_2\mathbf{b}_1+\mathbf{b}_2 = [1,2]$.

Constraints

• $W = W_2 W_1$ (matrix product) and $\mathbf{b} = W_2\mathbf{b}_1 + \mathbf{b}_2$.
• Return W of shape (o, i) and b of shape (o,).

Notes

• This is exactly why hidden layers need a nonlinear activation: insert a ReLU/sigmoid between the layers and this collapse no longer holds, so depth starts to add real expressive power.