Theorem. DGG head gradient and shared-weight occurrence sum [ftip-007Z]
✍️source

    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.

    Proof. 

      Softmax differentiation gives \(\nabla _{z_i}\mathcal L_i=E_i\). The outer-product rule gives the untied head identity. In the graph with independent copies, \(x_o\) does not depend on its own \(W_o\); all downstream paths from \(y_o\) are included in \(J_{io}\). The chain rule therefore gives \(\nabla _{W_o}\mathcal L_i=(J_{io}^{\mathsf T}E_i)x_o^{\mathsf T}\). Finally, the derivative of the diagonal map \(W_{\rm int}\mapsto (W_o=W_{\rm int})_{o\in \mathcal O_i}\) adds these partial derivatives.

      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.