Md. Asif Uddin
    I.6.X03

    The vector–Jacobian product, without forming the Jacobian

    symbolic▲▲△

    For y=xW\vec{y} = \vec{x}\mat{W} with x∈R1×n\vec{x} \in \R^{1 \times n} and W∈Rn×m\mat{W} \in \R^{n \times m}, write down the Jacobian ∂y/∂x\partial \vec{y}/\partial \vec{x} explicitly, count its entries, and then show that the quantity backpropagation actually needs — yˉJ\bar{\vec{y}}\mat{J} — can be computed without ever building it. State the cost of each route.

    Hint

    Ask what ∂yj/∂xi\partial y_j / \partial x_i is for this particular operation. The answer does not depend on x\vec{x} at all, which is why the product collapses.

    Solution

    The Jacobian. yj=∑ixiWijy_j = \sum_i x_i W_{ij}, so

    ∂yj∂xi=Wij\frac{\partial y_j}{\partial x_i} = W_{ij}

    The Jacobian is W\mat{W}, with nmnm entries. For a modest layer with n=m=1024n = m = 1024 that is 1 048 5761\,048\,576 numbers.

    The product. By Definition 7 the backward pass needs xˉ=yˉJT\bar{\vec{x}} = \bar{\vec{y}}\mat{J}^{\mathsf{T}}, and since J=W\mat{J} = \mat{W}:

    xˉ=yˉWT\bar{\vec{x}} = \bar{\vec{y}}\mat{W}^{\mathsf{T}}

    which is one matrix–vector product of exactly the shape the forward pass already performed. No new object is constructed.

    The costs. Forming J\mat{J} and multiplying costs O(nm)O(nm) storage and O(nm)O(nm) arithmetic. The VJP costs O(nm)O(nm) arithmetic and O(n+m)O(n + m) storage. The arithmetic is the same; the storage is not, and storage is what decides whether the layer fits.

    Why this generalises. The saving is not special to matrix products. For every primitive the derivative is known in closed form, and in every case the closed form is a rule for computing yˉJ\bar{\vec{y}}\mat{J} rather than a recipe for J\mat{J} — elementwise operations give a Hadamard product, a sum gives a broadcast, a broadcast gives a sum. A system that materialised Jacobians would be quadratic in width at every layer and could not train anything.

    Where it stops working. A second derivative needs the Jacobian of the Jacobian, and the same trick applies one level up — a Hessian–vector product is computable without the Hessian, by differentiating a VJP. But asking for the full Hessian of a model with NN parameters really does mean N2N^2 numbers, and no reformulation avoids it.

    Draws on