lgatr.layers.attention.config.SelfAttentionConfig

class lgatr.layers.attention.config.SelfAttentionConfig(in_mv_channels=None, out_mv_channels=None, in_s_channels=0, out_s_channels=0, additional_qk_mv_channels=0, additional_qk_s_channels=0, output_init='default', dropout_prob=None, num_heads=8, multi_query=False, attn_ratio=1, head_scale=False)[source]

Bases: object

Configuration for self-attention.

Parameters:

  • in_mv_channels (int | None) – Number of input multivector channels. Set automatically by the parent network.

  • out_mv_channels (int | None) – Number of output multivector channels. Set automatically by the parent network.

  • in_s_channels (int) – Input scalar channels. Use 0 for no scalar inputs. Set automatically by the parent network.

  • out_s_channels (int) – Output scalar channels. Use 0 for no scalar outputs. Set automatically by the parent network.

  • additional_qk_mv_channels (int) – Whether additional multivector features for the keys and queries will be provided. Set automatically by the parent network.

  • additional_qk_s_channels (int) – Whether additional scalar features for the keys and queries will be provided. Set automatically by the parent network.

  • output_init (str) – Initialization scheme for final linear layer. Set automatically by the parent network.

  • dropout_prob (float | None) – Dropout probability. Set automatically by the parent network.

  • num_heads (int) – Number of attention heads.

  • multi_query (bool) – Whether to do multi-query attention. Multi-query attention decreases memory consumption and parameter count by using a single set of keys and values for all heads.

  • attn_ratio (int) – Factor by which to increase the number of hidden channels (both multivectors and scalars). Vanilla transformers use 1; for multi-query 2 is more natural.

  • head_scale (bool) – Whether to use HeadScaleMHA following the NormFormer (https://arxiv.org/pdf/2110.09456). Each head is scaled by a learnable parameter before the heads are combined.

additional_qk_mv_channels: int = 0
additional_qk_s_channels: int = 0
attn_ratio: int = 1
classmethod cast(config)[source]

Cast a SelfAttentionConfig or mapping to a SelfAttentionConfig.

Return type:

SelfAttentionConfig

dropout_prob: float | None = None
head_scale: bool = False
property hidden_mv_channels: int

Number of hidden multivector channels.

property hidden_s_channels: int

Number of hidden scalar channels (0 if no scalar stream).

in_mv_channels: int | None = None
in_s_channels: int = 0
multi_query: bool = False
num_heads: int = 8
out_mv_channels: int | None = None
out_s_channels: int = 0
output_init: str = 'default'