law residual

bendtorch@1.3.0.0 · LAWS.bend:76

import Base
import bendtorch@1.3.0.0/main.bend as T
import bendtorch@1.3.0.0/spec.bend as Spec

76# LAW (residual): T.Residual's own code (T.ResidualOf: Seq(Fork(Id, f), Add2)), with a linear layer as f, gives a
77# backward pass that is the adjoint of its forward pass: for y = x + W·x and its dx, dy·y == dx·x. The backward pass
78# hands dy to both branches and adds what comes back, dx = dy + Wᵀ·dy: the gradient goes through the skip connection
79# untouched and through the layer, the rule every block of the Llama models relies on twice. (Over matrices the join
80# is Add2; here it is its twin over vectors, a + b forward and dy to both backward.)
81law residual:
82  for +w: List<&2, List<&2, Nat>>
83  for +x: List<&2, Nat>
84  for +dy: List<&2, Nat>
85  for g: List<&2, List<&2, Nat>>
86  {Spec.dotn(dy, Spec.fy(T.Two<T.Two<Unit, List<&2, List<&2, Nat>>>, Unit>, List<&2, Nat>, T.Two<T.Two<Unit, List<&2, Nat>>, Unit>, T.Layer.fwd(T.Two<T.Two<Unit, List<&2, List<&2, Nat>>>, Unit>, List<&2, Nat>, List<&2, Nat>, T.Two<T.Two<Unit, List<&2, Nat>>, Unit>, Spec.ResL(), T.Two{T.Two{Unit{}, w}, Unit{}}, x))) == Spec.dotn(Spec.bdx(T.Two<T.Two<Unit, List<&2, List<&2, Nat>>>, Unit>, List<&2, Nat>, T.Layer.bwd(T.Two<T.Two<Unit, List<&2, List<&2, Nat>>>, Unit>, List<&2, Nat>, List<&2, Nat>, T.Two<T.Two<Unit, List<&2, Nat>>, Unit>, Spec.ResL(), T.Two{T.Two{Unit{}, w}, Unit{}}, Spec.fsv(T.Two<T.Two<Unit, List<&2, List<&2, Nat>>>, Unit>, List<&2, Nat>, T.Two<T.Two<Unit, List<&2, Nat>>, Unit>, T.Layer.fwd(T.Two<T.Two<Unit, List<&2, List<&2, Nat>>>, Unit>, List<&2, Nat>, List<&2, Nat>, T.Two<T.Two<Unit, List<&2, Nat>>, Unit>, Spec.ResL(), T.Two{T.Two{Unit{}, w}, Unit{}}, x)), dy, T.Two{T.Two{Unit{}, g}, Unit{}})), x) : Nat}