lgatr.layers.attention.config.CrossAttentionConfig

class lgatr.layers.attention.config.CrossAttentionConfig(q_mv_channels=None, kv_mv_channels=None, out_mv_channels=None, out_s_channels=0, q_s_channels=0, kv_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 cross-attention.

Parameters:

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

  • kv_mv_channels (int | None) – Number of input key/value multivector channels. Set automatically by the parent network.

  • out_mv_channels (int | None) – Number of output multivector channels. 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.

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

  • kv_s_channels (int) – Input key/value scalar channels. Use 0 for no scalar inputs. 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.

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

Cast a CrossAttentionConfig or mapping to a CrossAttentionConfig.

Return type:

CrossAttentionConfig

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).

kv_mv_channels: int | None = None
kv_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'
q_mv_channels: int | None = None
q_s_channels: int = 0