LogScale

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

Bases: QueryScaling

Logarithmic query scaling.

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

\[\tilde{q}_{hi} = q_{hi} \cdot a_h \cdot \log(\max(n, 1)),\]

where \(h\) indexes attention heads and \(i\) indexes head channels.

Each attention head receives an unconstrained coefficient initialized independently and uniformly in [0, 1).

Parameters:

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

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

• dtype (dtype | None) – The dtype.