Every region is one matrix, and its rank is the count of live units

STATEMENT

Derive, in general, the affine map a two-layer ReLU network applies inside a single linear region. Then read two facts off the derivation: that the network is piecewise affine, and that the rank of the map on a region is bounded by the number of units switched on in it.

GIVEN

A network $\hat{\vec{y}} = \varphi(\vec{x}\mat{W}^{(1)} + \vec{b}^{(1)})\mat{W}^{(2)} + \vec{b}^{(2)}$ with $\varphi = \mathrm{ReLU}$, $\mat{W}^{(1)} \in \R^{d_0 \times d_1}$ and $\mat{W}^{(2)} \in \R^{d_1 \times d_2}$.

Fix an input $\vec{x}_\star$ and let $\mat{D} = \mathrm{diag}\!\left(\mathbf{1}[z^{(1)}_j > 0]\right)$ be the diagonal matrix of its activation pattern, with $k = \mathrm{tr}(\mat{D})$ units on.

FIND

The effective $\mat{W}_\star$ and $\vec{b}_\star$ with $\hat{\vec{y}} = \vec{x}\mat{W}_\star + \vec{b}_\star$ for every $\vec{x}$ in the region containing $\vec{x}_\star$, and a bound on $\operatorname{rank}(\mat{W}_\star)$.

STRATEGY

The one step that matters is replacing $\mathrm{ReLU}$ by $\mat{D}$. That substitution is exact — not an approximation — precisely on the set where the pattern does not change, and everything after it is matrix algebra.

SOLUTION

Step 1 — ReLU becomes a matrix. On the region, every unit’s sign is fixed. For a unit that is on, $\mathrm{ReLU}(z) = z$; for one that is off, $\mathrm{ReLU}(z) = 0$. Applying ReLU is therefore the same as multiplying on the right by $\mat{D}$:

$\varphi(\vec{z}) = \vec{z}\mat{D} \quad \text{for every } \vec{z} \text{ in the region}$

This is an identity on that set, not a linearisation.

Step 2 — substitute and expand.

which is equation (I.5.3). The network is affine on the region, with a $\mat{W}_\star$ and $\vec{b}_\star$ that depend on $\mat{D}$ and on nothing else.

Step 3 — piecewise affine. The regions are indexed by the finitely many patterns $\mat{D}$ that occur. Each contributes one affine map, and the maps agree on shared boundaries because a unit crossing zero contributes zero on both sides. So the network is a continuous piecewise affine function — a statement about the whole model, obtained without any assumption about the weights.

Step 4 — the rank bound. Rank cannot exceed that of any factor, and $\operatorname{rank}(\mat{D}) = k$:

using Rank, eigenvalues and the singular value decomposition · 0.LA.04. So on a region where only $k$ units are live, the network’s output lies in a subspace of dimension at most $k$, whatever the widths are elsewhere.

Step 5 — what that means. Width $d_1$ does not act on a single input. On any one input only the live units count, and the network behaves like a much smaller model. Width buys variety across regions, not capacity within one.

Check — sanity

It reproduces I.5.B01. There $d_0 = 2$, $d_1 = 3$, $d_2 = 1$, $\mat{D} = \mathrm{diag}(1,1,0)$ and $k = 2$. The formula gives $\vec{w}_\star = (3, 0)$ and $b_\star = -0.5$, which is what the direct forward pass returned.

The rank bound is not vacuous there. $\min(2, 2, 1) = 1$, and $\vec{w}_\star = (3,0)$ has rank 1. The bound is tight in that case.

Every ReLU disappears. If any survived, the expression would not be affine, and the substitution in Step 1 would have been wrong. Counting them is a real check: three units in, three columns of $\mat{D}$, none left over.