Gated Attention

Arcee Trinity

Medium · 🔒 Plus required

Gate projection on attention output

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

Trinity attention is grouped-query attention plus QK-norm, then an extra sigmoid gate on the attention output before W^O. Section 2.2 cites Qiu et al. (2025): the gate is there to shrink attention sinks, clip huge activations, help long-context generalization, and cut loss spikes — the same stability theme as QK-norm and sandwich RMSNorm.

The gate is not a router and not an attention-score mask. It is a per-token, per-channel multiplier produced from the layer input by a full-width linear map. Every layer, local or global, dense or MoE, uses the same gate.

How it works

After QK-norm and (on local layers) RoPE, head i produces the usual SDPA mix \mathbf{o}^{\mathrm{sdpa}}_{t,i}\in\mathbb{R}^{d_h}. Equations (11)–(14) then gate and project:

\mathbf{g}_{t}=\sigma(W^{G}\mathbf{x}_{t}),\qquad \mathbf{g}_{t,i}=\mathrm{split}_{h_q}(\mathbf{g}_{t})_{i},

\widetilde{\mathbf{o}}_{t,i}=\mathbf{o}^{\mathrm{sdpa}}_{t,i}\odot\mathbf{g}_{t,i},\qquad \mathbf{u}_{t}=W^{O}\,[\widetilde{\mathbf{o}}_{t,1};\ldots;\widetilde{\mathbf{o}}_{t,h_q}].

Here \sigma is elementwise sigmoid, W^{G},W^{O}\in\mathbb{R}^{d\times d} (no bias), and \mathrm{split}_{h_q} cuts \mathbf{g}_t\in\mathbb{R}^{d} into h_q contiguous d_h-vectors. On Large, d=3072, h_q=48, d_h=128, so the split is exactly the head axis of the concatenated SDPA output.

\mathbf{x}_t in §2.2 is the attention-sublayer input. In the sandwich block that is the pre-normed residual (input_layernorm output), not the raw skip. Official AfmoeAttention computes gate_states = self.gate_proj(hidden_states) from that same tensor, then after SDPA does output * torch.sigmoid(gate_states) and o_proj. Sigmoid is applied to the projected gate, matching \sigma(W^G\mathbf{x}_t) rather than gating in logit space.

Values are not gated and not QK-normalized. Only Q and K go through RMSNorm; V is a plain projection. The gate never enters the softmax.

Because the gate is a function of the destination token only, it can suppress a whole head’s mixed value at that position even when the softmax is sharp. That is the intended sink/activation control: a position that would otherwise dump a large \mathbf{o}^{\mathrm{sdpa}} can be driven toward zero by \sigma(W^G\mathbf{x}_t)\approx 0 on those channels.

Official code

The report repo has no modeling source. The gate lives in AfmoeAttention in modeling_afmoe.py (gate_proj, then output * torch.sigmoid(gate_states) before o_proj). modular_afmoe.py is the file editors are told to change; modeling_afmoe.py is generated from it.

Watch-outs

• Gate after SDPA, before W^O. Gating the residual, the pre-norm input, or the logits is a different module than Qiu/Trinity.
• Apply sigmoid to W^G\mathbf{x}, not to \mathbf{x}. The paper’s \sigma is on the projection.
• W^G is d\to h_q d_h. On Large that equals d\to d, but the last axis must align with concatenated heads. A d\to h_{kv}d_h map (KV width) will not broadcast.
• The gate uses the attention input \mathbf{x}_t, not Q or the SDPA output. Feeding \mathbf{o}^{\mathrm{sdpa}} into W^G makes a recurrent-looking gate the paper does not define.

Sources

• Paper: Arcee Trinity Large Technical Report, §2.2 equations (11)–(14)
• Code: arcee-ai/trinity-large-tech-report; transformers AfmoeAttention