QASSMax#
- class sdm.nn.QASSMax(channels: int, num_heads: int, hidden_channels: int = 64, device: device | str | None = None, dtype: dtype | None = None)#
Bases:
QueryScalingQuery-Aware Scalable SoftMax (QASSMax).
This scaling method was introduced in the “TabICLv2: A better, faster, scalable, and open tabular foundation model” paper as a temperature-scaling method for attention.
For a query tensor \(q\) and key length \(n\), this module returns a scaled query
\[\tilde{q}_{hi} = q_{hi} \cdot \mathrm{MLP}_{\mathrm{base}}(\log n)_{hi} \cdot \left(1 + \tanh(\mathrm{MLP}_{\mathrm{gate}}(q_h)_i)\right),\]where \(h\) indexes attention heads, \(i\) indexes head channels, \(\mathrm{MLP}_{\mathrm{base}}\) maps the log key length to per-head, per-channel scale factors, and \(\mathrm{MLP}_{\mathrm{gate}}\) maps each per-head query vector to a bounded query-dependent gate.
Multiplying the query scales the subsequent attention logits while keeping the attention computation compatible with standard softmax kernels. The length-dependent factor counteracts attention fading as the number of keys grows, while the query-dependent gate lets the scale vary across queries.