One conformability error in five stages

STATEMENT

A five-stage pipeline is specified below. Exactly one stage cannot be computed. Find it, say why, and state the shape the offending tensor must have for the pipeline to run.

GIVEN

A batch of $B = 32$ examples, each a flattened $28\times 28$ image:

Elementwise nonlinearities sit between the stages. They do not change shape, so they can be ignored for this question — which is itself worth noticing.

FIND

The failing stage, the reason, and the corrected shape of the offending weight matrix. Then the total parameter count once it is corrected, counting weights only.

STRATEGY

Propagate the shape forward one stage at a time. A matrix product needs the inner dimensions to agree, so at each stage compare the running width with the first dimension of the next weight matrix; the first disagreement is the error.

SOLUTION

Step 1 — propagate. Write the running shape after each stage. The rule is $(a \times b)(b \times c) \to (a \times c)$: the inner pair must match and then vanishes, while the outer pair survives.

Step 2 — name the failure. At stage 3 the running width is $128$, because stage 2 projected down to $128$. But $\mat{W}^{(3)}$ expects an input of width $256$. The inner dimensions are $128$ and $256$; they disagree, so the product is undefined.

Note the batch axis played no part. It is carried along untouched by every stage, which is why a shape error in a stack of dense layers is always a statement about widths and never about how many examples you fed in.

Step 3 — correct it. A weight matrix’s first dimension is the width it consumes and its second is the width it produces. The width consumed must be $128$; the width produced was intended to be $64$ and nothing contradicts that. So

$\mat{W}^{(3)} : 128 \times 64$

and the pipeline then runs $32\times784 \to 256 \to 128 \to 64 \to 10$.

Step 4 — count the parameters. Weights only, so the sum of the products:

Check — sanity

The first layer dominates, as it must. $200{,}704$ of $242{,}304$ weights — $82.8\%$ — sit in the first matrix, because it is the only one touching the $784$-dimensional input. Any parameter count for a network with a wide input and a narrowing stack should be dominated by its first layer; if yours is not, recheck.

The corrected chain telescopes. Reading the widths in order gives $784 \to 256 \to 128 \to 64 \to 10$, where each weight matrix’s second dimension is the next one’s first. That chain property is the shape rule restated, and it is the fastest way to eyeball a whole architecture.