law gqa_head

bendtorch@1.3.0.0 · LAWS.bend:239

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

239# LAW (gqa_head): with g = 1 + i + k query heads per key/value head, query head j·g + i reads key/value head
240# (j·g + i) / g = j, the head of its own group: T.Attention's forward pass picks head hi / g, and its backward pass
241# sums the gradients of heads j·g + 0 .. j·g + g - 1 into key/value head j, so both see the same groups.
242law gqa_head:
243  for +j: Nat
244  for +i: Nat
245  for +k: Nat
246  {Nat.div(Nat.add(Nat.mul(j, 1n+Nat.add(i, k)), i), 1n+Nat.add(i, k)) == j : Nat}