Derive Backpropagation for Matrix-Product Layers

Think of a neural-network block that forms its output by first multiplying a sequence of trainable weight matrices, then applying that composite matrix to an input. The block has no nonlinearity inside it — it is purely a chain of matrix multiplications.

Denote the trainable matrices by $W_1, W_2, \ldots, W_n$. Their product defines one effective linear map

$$C = W_1 W_2 \cdots W_n.$$

For an input vector or mini-batch $X$, the forward computation is

$$Z = C X = W_1 W_2 \cdots W_n X.$$

A scalar loss $\mathcal{L}$ is given, and the upstream gradient

$$G = \frac{\partial \mathcal{L}}{\partial Z}$$

is provided by the loss or by downstream layers (so $G$ has the same shape as $Z$).

Derive the backward pass for this block. In particular:

1. Give the gradient $\frac{\partial \mathcal{L}}{\partial W_j}$ for each individual matrix $W_j$, for all $1 \le j \le n$.
2. Show explicitly how the multivariate chain rule (or matrix differential) applies to the matrix product — that is, why the answer takes the form it does, not merely the final formula.
3. Check that the resulting gradient $\frac{\partial \mathcal{L}}{\partial W_j}$ has the same shape as $W_j$.
4. Describe an efficient implementation that does not recompute the same prefix and suffix matrix products for every $j$.

Hint: Isolate one matrix. $W_j$ sits in the middle of the product, which makes the chain rule awkward. Try holding every other matrix fixed and viewing $Z$ as a composition of three pieces: a left prefix formed by the matrices before $W_j$, the target $W_j$ itself, and a right suffix formed by the matrices after it (the suffix may absorb $X$). Relative to $W_j$ alone, what kind of map is this — and what does that give you?

Hint: Differentiate the isolated form. Once $W_j$ is isolated, perturb only $W_j$ and ask how $Z$ changes to first order. For a scalar loss, the contraction of the upstream gradient with a matrix change is naturally written as a trace: $d\mathcal{L} = \operatorname{tr}(G^\top dZ)$. The cyclic invariance of the trace is what lets you gather everything multiplying $dW_j$ into a single factor. Be careful about which side each frozen piece ends up on and where transposes appear.

Hint: Avoid the $O(n^2)$ trap. The per-matrix formula reuses overlapping prefix/suffix products, so recomputing them for every $j$ is quadratic. Ask: do you ever need the prefix as a full matrix, or only its action on the upstream gradient? And can all suffix pieces be obtained from a single pass in one direction? Two coordinated sweeps in opposite directions should be enough to produce every gradient in linear time.

Constraints & Assumptions

• Shapes: let $W_j \in \mathbb{R}^{d_{j-1} \times d_j}$, so the composite $C \in \mathbb{R}^{d_0 \times d_n}$. For a mini-batch, $X \in \mathbb{R}^{d_n \times B}$ and $Z, G \in \mathbb{R}^{d_0 \times B}$; for a single example take $B = 1$.
• The matrices need not be square; adjacent inner dimensions only have to match so the product is defined.
• There is no nonlinearity or bias term inside this block — it is a pure chain of matrix multiplications. Any activation lives in the layers that produce $G$.
• "Efficient" means linear in the number of matrices $n$: the whole backward pass should cost $O(n)$ matrix products, not $O(n^2)$, and should avoid materializing large $d_0 \times d_{j-1}$ prefix matrices when $B$ is small.

Clarifying Questions to Ask

• Is $X$ a single vector or a mini-batch, and is the loss summed or averaged over the batch (which scales every $dW_j$)?
• Should I treat $X$ as fixed input, or is the gradient $\partial\mathcal{L}/\partial X$ also required?
• Are the matrices guaranteed non-singular / square, or fully general rectangular?
• Is there any nonlinearity or bias term between the matrices, or is the block a pure linear chain as stated?
• What is the target deliverable — closed-form gradient formula, an $O(n)$ algorithm, or working code/pseudocode?
• Do we ever need to materialize the composite $C$ at inference, or only its action $CX$?

What a Strong Answer Covers

• Cleanly isolates a single $W_j$ by grouping the rest of the product into a left prefix and a right suffix, and states the empty-product-is-identity convention.
• A genuine derivation (matrix differential + trace cyclic property, or an equivalent index/chain-rule argument) that produces the per-matrix gradient rather than asserting it, with transposes and left/right placement justified.
• A closed-form expression for $\partial\mathcal{L}/\partial W_j$ in terms of the prefix, the upstream gradient $G$, and the suffix.
• An explicit dimension check confirming the gradient has the same shape as $W_j$ (i.e. $d_{j-1}\times d_j$), with the batch dimension correctly contracted away.
• A linear-time ($O(n)$) algorithm that computes all $n$ gradients without recomputing overlapping products or materializing the large prefix matrices — i.e. coordinated sweeps in both directions, with correct transpose placement.
• Awareness of practical concerns: batch averaging applied exactly once, and the link to standard reverse-mode autodiff over a linear chain.

Follow-up Questions

• How do the cost and the formula change if you also need $\partial\mathcal{L}/\partial X$?
• Numerically, why is composing the matrices into a single $C$ and backpropagating through $C$ alone not equivalent — what information about the individual $W_j$ would you lose?
• If $n$ is large, what conditioning / vanishing-or-exploding-gradient issues arise from the repeated products, and how would you mitigate them?
• Suppose two of the matrices are tied (the same parameter reused at positions $j$ and $k$). How does the gradient with respect to that shared parameter change?

Overview: This question tests understanding of backpropagation through linear blocks, matrix calculus, and computing gradients with respect to individual weight matrices in a product of trainable matrices, including reasoning about tensor shapes and matrix differentials.