Theorem. DGG head gradient and shared-weight occurrence sum [ftip-007Z]
AGENTDRAFTED
Fix a sampled history and an active token \(i\) in the interior of the
unclipped branch of the GRPO objective. Consider a finite directed acyclic
computation graph that is differentiable at the parameter point in question.
Let \(a_i\) be the sampled token, \(\widehat A_i\) its fixed normalized advantage, and
\(p_i=\operatorname {softmax}(z_i)\) the current distribution on a finite
vocabulary \(\mathcal V\). With fixed behavior probability \(b_i>0\), set
\[
r_i=\frac {p_i(a_i)}{b_i},
\qquad
\mathcal L_i=r_i\widehat A_i,
\qquad
E_i=r_i\widehat A_i(e_{a_i}-p_i).
\]
Here \(e_{a_i}\in \mathbb R^{|\mathcal V|}\) is the sampled-token basis vector.
The history and sampling decisions are held fixed during differentiation.
Suppose the output head \(W_{\rm lm}\in
\mathbb R^{|\mathcal V|\times d_{\rm model}}\) is untied: its only path to
\(\mathcal L_i\) is through \(z_i=W_{\rm lm}h_{L,i}\), and
\(h_{L,i}\in \mathbb R^{d_{\rm model}}\) is independent of \(W_{\rm lm}\).
Then
\[
G_i^{\rm lm}=\nabla _{W_{\rm lm}}\mathcal L_i
=E_i h_{L,i}^{\mathsf T}.
\]
Let \(W_{\rm int}\in \mathbb R^{m\times d}\) be a shared intermediate
weight. Index by a finite set \(\mathcal O_i\) every occurrence of this weight
that can affect \(z_i\), assuming that all such uses have the form
\(y_o=W_{\rm int}x_o\), with \(x_o\in \mathbb R^d\). First replace these uses
by independent copies \(W_o\) and evaluate them all at \(W_o=W_{\rm int}\).
Let \(J_{io}\in \mathbb R^{|\mathcal V|\times m}\) be the downstream Jacobian
from node \(y_o\) to \(z_i\) in this graph, with the other weight copies fixed.
Define the contribution of occurrence \(o\) by
\[
H_{io}=\nabla _{W_o}\mathcal L_i
=(J_{io}^{\mathsf T}E_i)x_o^{\mathsf T}.
\]
On tying the copies, the total derivative is
\[
G_i^{\rm int}=\nabla _{W_{\rm int}}\mathcal L_i
=\sum _{o\in \mathcal O_i}H_{io}.
\]
The head gradient and each occurrence contribution have rank at most one;
the total shared-weight gradient need not.
For a layer used once per position in a causal network, the sum includes
every earlier position whose output affects token \(i\). Reuse across depth
adds further occurrences. Tying the head to embeddings or other blocks also
requires their contributions; the displayed head identity assumes that such
tying is absent.
Clipped or inactive terms require their own derivative or mask.
[miao2026when, Section 4.2.1, Proposition 1, and Appendix A.1]
derives the intermediate outer product through one local application. For a
shared weight, that calculation gives \(H_{io}\); identifying it with the total
\(G_i^{\rm int}\) omits the other occurrences. The sum above supplies the
chain rule needed for shared parameters.
