The Jacobian of a residual stream is an ordered product

For a residual stream h_{ℓ+1} = h_ℓ + w_ℓ(h_ℓ) with differentiable layers w_ℓ on a real normed space, the derivative of the whole stream is the ORDERED product of the layer Jacobians I + w_ℓ′, each evaluated at the point the stream has reached (Theorem 1 of transformers-attend-III, «recurrence factorisation»). Expanding the product gives J − I = Σ_k w_k′ + (cross terms).

Main results

• step, flow, jac: one residual layer, the stream through a list of layers (head first), and the ordered product of I + fderiv w along the trajectory.
• hasFDerivAt_flow: HasFDerivAt (flow ws) (jac ws x) x.

References

• C. Marín, Transformers attend (Part III), Theorem 1 (thm:bridge).

def TransformerAttention.ResidualJacobian.step {E : Type u_1} [NormedAddCommGroup E] (w : E → E) : E → E

One residual layer h ↦ h + w h.

Equations

• TransformerAttention.ResidualJacobian.step w h = h + w h

def TransformerAttention.ResidualJacobian.flow {E : Type u_1} [NormedAddCommGroup E] : List (E → E) → E → E

The residual stream through the layers of a list, the head first.

Equations

• TransformerAttention.ResidualJacobian.flow [] = id
• TransformerAttention.ResidualJacobian.flow (w :: ws) = TransformerAttention.ResidualJacobian.flow ws ∘ TransformerAttention.ResidualJacobian.step w

noncomputable def TransformerAttention.ResidualJacobian.jac {E : Type u_1} [NormedAddCommGroup E] [NormedSpace ℝ E] : List (E → E) → E → E →L[ℝ] E

The ordered product of the layer Jacobians I + w′ along the trajectory, the head layer first (rightmost).

Equations

• One or more equations did not get rendered due to their size.
• TransformerAttention.ResidualJacobian.jac [] x✝ = ContinuousLinearMap.id ℝ E

theorem TransformerAttention.ResidualJacobian.hasFDerivAt_flow {E : Type u_1} [NormedAddCommGroup E] [NormedSpace ℝ E] (ws : List (E → E)) (hw : ∀ w ∈ ws, Differentiable ℝ w) (x : E) : HasFDerivAt (flow ws) (jac ws x) x

Recurrence factorisation: for differentiable layers, the derivative of the residual stream at x is the ordered product of the layer Jacobians I + w_k′ evaluated along the trajectory from x.