flâneur

Softplus Gated Attention > Sigmoid Gated Attention

1a3orn.com · 1,032 words · saved by 1 readers

<h1>TLDR</h1> <p>Using <a href="https://en.wikipedia.org/wiki/Softplus">softplus</a> rather than sigmoid as an activation function for the gate in a <a href="https://arxiv.org/pdf/2505.06708">gated attention</a> Transformer notably improves performance, just as gated attention itself seems to improve performance over vanilla attention.</p> <p>In my small testbed, softplus is ~10% to ~30% faster than sigmoid at reaching equal loss, or ~25% to ~50% faster than non-gated attention, with the effect possibly (?) increasing with size.</p> <p>The effect seems robust across every size I've tested, although as a GPU-poor person I've only gone up to 2 billion tokens and 95 million parameters.</p> <h1>Architecture</h1> <p>The Qwen team recently <a href="https://arxiv.org/pdf/2505.06708">published</a> a study of gated attention.</p> <p>In the mainline dense implementation (at "G1" in the paper, with a sigmoid), this introduces a multiplicative per-channel rescaling of the concatenated outputs of MHA, where each channel is scaled according to a gate learned from the per-token input to the attention.</p> <p><img src="/gated_01.png" alt="Qwen 3"></p> <p>This requires a new <code>dim_model x (num_head x dim_head)</code> parameter linear layer -- in practice just <code>dim_model^2</code> -- that per-token and per-channel modulates the output of MHA.</p> <p>This is simple to implement -- here's it added to <a href="https://github.com/karpathy/nanoGPT">nanoGPT</a>:</p> <pre><code class="language-python">class CausalSelfAttention(nn.Module): def __init__(self, config): super().__init__() assert config.n_embd % config.n_head == 0 self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias) self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias) # New! self.gate = nn.Linear(config.n_embd, config.n_embd, bias=False) # etc etc etc def forward(self, original_input): # ... get results fo

TLDR Using softplus rather than sigmoid as an activation function for the gate in a gated attention Transformer notably improves performance, just as gated attention itself seems to improve performance over vanilla attention. In my small testbed, softplus is ~10% to ~30% faster than sigmoid at reaching equal loss, or ~25% to ~50% faster than non-gated attention, with the effect possibly (?) increasing with size. The effect seems robust across every size I've tested, although as a GPU-poor person I've only gone up to 2 billion tokens and 95 million parameters. Architecture The Qwen team…

saved by

related reading