GatedLogScale

class sdm.nn.GatedLogScale(channels: int, num_heads: int, hidden_channels: int, device: device | str | None = None, dtype: dtype | None = None)

Bases: LogScale

Gated logarithmic query scaling.

For a query tensor \(q\), effective key length \(n\), learned per-head coefficient \(a_h\), and learned gate \(g\), this module returns

\[\tilde{q}_{hi} = q_{hi} \cdot a_h \cdot \log(\max(n, 1)) \cdot \left(1 + \tanh(g(q_h)_i)\right).\]

The multiplicative gate is initialized to one and bounded between zero and two. Per-head coefficients are initialized independently and uniformly in [0, 1).

Parameters:

• channels (int) – The number of channels per attention head.

• num_heads (int) – The number of query attention heads.

• hidden_channels (int) – The hidden width of the gate MLP.

• device (device | str | None) – The device.

• dtype (dtype | None) – The dtype.