dual_attention Package Reference

dual_attention.attention module

An implementation of attention including several additional features and customizations over the standard pytorch implementation.

class dual_attention.attention.Attention(d_model: int, n_heads: int, dropout: float, key_dim: int = None, n_kv_heads: int = None, add_bias_kv: bool = False, add_bias_out: bool = False, total_n_heads: int = None)

Bases: Module

__init__(d_model: int, n_heads: int, dropout: float, key_dim: int = None, n_kv_heads: int = None, add_bias_kv: bool = False, add_bias_out: bool = False, total_n_heads: int = None)

An implementation of Attention with some added customization.

Allows multi-query attention/grouped query attention, rotary positional embeddings, and custom relation activation functions.

Parameters:

• d_model (int) – model dimension

• n_heads (int) – number of heads (query heads if n_kv_heads is set)

• dropout (float) – dropout rate

• n_kv_heads (int, optional) – number of key/value heads. used to implement multi-query attention or grouped query attention. n_kv_heads=1 corresponds to MQA, n_kv_heads > 1 corresponsd to grouped query attention. n_kv_heads=n_heads is standard MHA. uses MHA when None. By default None

• add_bias_kv (bool, optional) – whether to use bias in key/value projections, by default False

• add_bias_out (bool, optional) – whether to use bias in out projection, by default False

• total_n_heads (int, optional) – total number of heads in dual attention (if using dual attention). used to ensure that concat(A, E) is of dimension d_model after concatentation. hence, output dimension is (d_model // total_heads) * n_heads. if None, total_heads = n_heads and output dimension is d_model

forward(query: Tensor, key: Tensor, value: Tensor, freqs_cos: Tensor = None, freqs_sin: Tensor = None, attn_mask: Tensor = None, is_causal: bool = False, need_weights: bool = False)

compute attention with given query, key, value.

if freqs_cos and freqs_sin are given, apply rotary positional embeddings. if attn_mask is given, apply attention mask. if is_causal is True, apply causal mask (attn_mask must be None).

Parameters:

• query (torch.Tensor) – query sequence of shape [bsz, len_in, d_model]

• key (torch.Tensor) – key sequence of shape [bsz, len_ctx, d_model]

• value (torch.Tensor) – value sequence of shape [bsz, len_ctx, d_model]

• freqs_cos (torch.Tensor, optional) – cosine of frequencies for RoPE. RoPE is applied if given. Note: RoPE does not work for cross-attention. By default None

• freqs_sin (torch.Tensor, optional) – cosine of frequencies for RoPE. RoPE is applied if given. Note: RoPE does not work for cross-attention. By default None

• attn_mask (torch.Tensor, optional) – boolean attention mask of shape [len_in, len_ctx]. True at [i,j] indicates i is allowed to attend to j. By default None

• is_causal (bool, optional) – whether to apply a causal mask. If True, attn_mask must be None. Only applies for self-attention. By default False

• need_weights (bool, optional) – whether to return the attention scores. If True, return value will be tuple (output, attn_scores). If True, will compute attention manually rather than using flash attention. By default False

Returns:

result of attention

Return type:

torch.Tensor

dual_attention.attention_utils module

dual_attention.attention_utils.repeat_kv(x: Tensor, n_rep: int) → Tensor

torch.repeat_interleave(x, dim=2, repeats=n_rep)

dual_attention.attention_utils.precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0)

dual_attention.attention_utils.reshape_for_broadcast(freqs_cis: Tensor, x: Tensor)

dual_attention.attention_utils.apply_rotary_emb(xq: Tensor, xk: Tensor, freqs_cos: Tensor, freqs_sin: Tensor) → Tuple[Tensor, Tensor]

dual_attention.attention_utils.compute_diag_mask(size, device=None)

computes an attention mask with False on the diagonal and True elsewhere

dual_attention.attention_utils.compute_causal_mask(size, device=None)

computes an attention mask with True at (i,j) if i <= j

dual_attention.dual_attention module

This module implements Dual Attention: a variant of multi-head attention with two distinct types of attention heads: self-attention and relational attention.

class dual_attention.dual_attention.DualAttention(d_model: int, n_heads_sa: int, n_heads_ra: int, dropout: float, sa_kwargs: dict = None, ra_kwargs: dict = None, share_attn_params: bool = False, ra_type: str = 'relational_attention')

Bases: Module

__init__(d_model: int, n_heads_sa: int, n_heads_ra: int, dropout: float, sa_kwargs: dict = None, ra_kwargs: dict = None, share_attn_params: bool = False, ra_type: str = 'relational_attention')

An implementation of Dual Attention.

The DualAttention module is a form of multi-head attention involving a composition of two distinct types of attention heads. The first type is standard self-attention, which captures object-level (i.e., sensory) features, and the second type is relational attention, which captures relational features.

Parameters:

• d_model (int) – model dimension

• n_heads_sa (int) – number of self-attention heads

• n_heads_ra (int) – number of relational attention heads

• dropout (float) – dropout rate

• sa_kwargs (dict, optional) – self-attention kwargs, by default None

• ra_kwargs (dict, optional) – relational attention kwargs, by default None

• share_attn_params (bool, optional) – whether to share attention parameters between self-attention and relational attention. If True, w{q,k} in sensory attention and w{q,k}_attn in relational attention are shared. number of heads in each must be the same. By default False

• ra_type (str, optional) – type of relational attention module (e.g., whether to use RCA for an ablation experiment). by default ‘relational_attention’.

forward(x: Tensor, symbols: Tensor, attn_mask: Tensor = None, is_causal: bool = False, freqs_cos: Tensor = None, freqs_sin: Tensor = None, need_weights: bool = False)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

dual_attention.dual_attn_blocks module

This module implements Encoder and Decoder Blocks for a Dual Attention Transformer. Each block consists of Dual Attention, (Cross-Attention in the case of the Decoder), Feed-Forward Block, LayerNorms/Residuals.

class dual_attention.dual_attn_blocks.DualAttnEncoderBlock(d_model: int, n_heads_sa: int, n_heads_ra: int, dff: int, activation: str, dropout_rate: float, norm_first: bool, norm_type: str = 'layernorm', sa_kwargs: dict = None, ra_kwargs: dict = None, ra_type: str = 'relational_attention', share_attn_params: bool = False, bias: bool = True, causal: bool = False)

Bases: Module

__init__(d_model: int, n_heads_sa: int, n_heads_ra: int, dff: int, activation: str, dropout_rate: float, norm_first: bool, norm_type: str = 'layernorm', sa_kwargs: dict = None, ra_kwargs: dict = None, ra_type: str = 'relational_attention', share_attn_params: bool = False, bias: bool = True, causal: bool = False)

Dual Attention Encoder Block.

A Dual Attention Encoder is a variant of the Transformer Encoder that uses a combination of two distinct types of attention heads. The first type is standard self-attention, which captures object-level (i.e., sensory) features, and the second type is relational attention, which captures relational features.

Parameters:

• d_model (int) – model dimension.

• n_heads_sa (int) – number of standard self-attention heads.

• n_heads_ra (int) – number of relational attention heads.

• dff (int) – intermediate dimension of feed-forward block.

• activation (str) – name of activation function to use in feedforward block.

• dropout_rate (float) – dropout rate.

• norm_first (bool) – whether to apply normalization before or after attention. norm_first=True means pre-norm otherwise post-norm.

• norm_type ('layernorm' or 'rmsnorm, optional) – type of normalization to use, by default ‘layernorm’

• sa_kwargs (dict, optional) – self-attention kwargs, by default None

• ra_kwargs (dict, optional) – relational attention kwargs, by default None

• ra_type (str, optional) – type of relational attention module (e.g., whether to use RCA for an ablation experiment), by default ‘relational_attention’

• share_attn_params (bool, optional) – whether to share attention parameters between self-attention and relational attention. If True, w{q,k} in sensory attention and w{q,k}_attn in relational attention are shared. number of heads in each must be the same. By default False

• bias (bool, optional) – whether to use bias in multi-head attention, by default True

• causal (bool, optional) – whether attention operations should be causal, by default False

forward(x, symbols, freqs_cos=None, freqs_sin=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dual_attention.dual_attn_blocks.DualAttnDecoderBlock(d_model: int, n_heads_sa: int, n_heads_ra: int, n_heads_cross: int, dff: int, activation: str, dropout_rate: float, norm_first: bool, norm_type: str = 'layernorm', sa_kwargs: dict = None, ra_kwargs: dict = None, cross_kwargs: dict = None, ra_type: str = 'relational_attention', share_attn_params: bool = False, bias: bool = True, causal: bool = True)

Bases: Module

__init__(d_model: int, n_heads_sa: int, n_heads_ra: int, n_heads_cross: int, dff: int, activation: str, dropout_rate: float, norm_first: bool, norm_type: str = 'layernorm', sa_kwargs: dict = None, ra_kwargs: dict = None, cross_kwargs: dict = None, ra_type: str = 'relational_attention', share_attn_params: bool = False, bias: bool = True, causal: bool = True)

Dual Attention Decoder Block.

A Dual Attention Decoder is a variant of the Transformer Decoder that uses a combination of two distinct types of attention heads. The first type is standard self-attention, which captures object-level (i.e., sensory) features, and the second type is relational attention, which captures relational features.

Parameters:

• d_model (int) – model dimension.

• n_heads_sa (int) – number of standard self-attention heads.

• n_heads_ra (int) – number of relational attention heads.

• n_heads_cross (int) – number of cross-attention heads.

• dff (int) – intermediate dimension of feed-forward block.

• activation (str) – name of activation function to use in feedforward block.

• dropout_rate (float) – dropout rate.

• norm_first (bool) – whether to apply normalization before or after attention. norm_first=True means pre-norm otherwise post-norm.

• norm_type ('layernorm' or 'rmsnorm, optional) – type of normalization to use, by default ‘layernorm’

• sa_kwargs (dict, optional) – self-attention kwargs, by default None

• ra_kwargs (dict, optional) – relational attention kwargs, by default None

• cross_kwargs (dict, optional) – cross-attention kwargs, by default None

• ra_type (str, optional) – type of relational attention module (e.g., whether to use RCA for an ablation experiment), by default ‘relational_attention’

• share_attn_params (bool, optional) – whether to share attention parameters between self-attention and relational attention. If True, w{q,k} in sensory attention and w{q,k}_attn in relational attention are shared. number of heads in each must be the same. By default False

• bias (bool, optional) – whether to use bias in multi-head attention, by default True

• causal (bool, optional) – whether attention operations should be causal, by default False

forward(x, context, symbols)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

dual_attention.hf module

This module implements a Huggingface Hub mixin for the Dual Attention Transformer Language Model.

class dual_attention.hf.DualAttnTransformerLM_HFHub(*args, **kwargs)

Bases: PyTorchModelHubMixin, DualAttnTransformerLM

Huggingface Hub mixin for the Dual Attention Transformer Language Model.

This class inherits from both the PyTorchModelHubMixin and the DualAttnTransformerLM classes. It allows you to load a DAT model from Huggingface Hub using the from_pretrained method.

Example: ```python from dual_attention.hf import DualAttnTransformerLM_HFHub

model = DualAttnTransformerLM_HFHub.from_pretrained(“awni00/DAT-sa8-ra8-ns1024-sh8-nkvh4-343M”) ```

__init__(*args, **kwargs)

class dual_attention.hf.TransformerLM_HFHub(*args, **kwargs)

Bases: PyTorchModelHubMixin, TransformerLM

Huggingface Hub mixin for the Transformer Language Model.

Used for the purposes of loading baselines trained via the same procedure as DAT models.

This class inherits from both the PyTorchModelHubMixin and the TransformerLM classes. It allows you to load a DAT model from Huggingface Hub using the from_pretrained method.

Example: ```python from dual_attention.hf import TransformerLM_HFHub

model = TransformerLM_HFHub.from_pretrained(“awni00/T-sa24-757M”) ```

__init__(*args, **kwargs)

dual_attention.language_models module

This module implements the Dual Attention Transformer Language Model (and a standard Transformer Language Model as a baseline).

The DAT-LM uses a “Decoder-only” architecture, causally processing input tokens and predicting the next token in the sequence.

class dual_attention.language_models.DualAttnTransformerLM(vocab_size: int, d_model: int, n_layers: int, n_heads_sa: int, n_heads_ra: int, symbol_retrieval_kwargs: dict, dff: int, dropout_rate: float, activation: str, norm_first: bool, max_block_size: int, norm_type: str = 'layernorm', sa_kwargs: dict = None, ra_kwargs: dict = None, ra_type: str = 'relational_attention', share_attn_params: bool = False, symbol_retrieval: str = 'symbolic_attention', symbol_retriever_config: dict = None, pos_enc_type: str = 'pos_emb', bias: bool = True)

Bases: Module

Dual Attention Transformer Language Model

__init__(vocab_size: int, d_model: int, n_layers: int, n_heads_sa: int, n_heads_ra: int, symbol_retrieval_kwargs: dict, dff: int, dropout_rate: float, activation: str, norm_first: bool, max_block_size: int, norm_type: str = 'layernorm', sa_kwargs: dict = None, ra_kwargs: dict = None, ra_type: str = 'relational_attention', share_attn_params: bool = False, symbol_retrieval: str = 'symbolic_attention', symbol_retriever_config: dict = None, pos_enc_type: str = 'pos_emb', bias: bool = True)

Dual Attention Transformer Language Model.

Parameters:

• vocab_size (int) – vocabulary size.

• d_model (int) – model dimension.

• n_layers (int) – number of layers.

• n_heads_sa (int) – number of self-attention heads in dual-attention.

• n_heads_ra (int) – number of relational attention heads in dual-attention.

• symbol_retrieval_kwargs (dict) – keyword arguments for symbol retrieval module.

• dff (int) – size of intermediate layer in feedforward blocks.

• dropout_rate (float) – dropout rate.

• activation (str) – name of activation function (e.g., ‘relu’, ‘gelu’, or ‘swiglu’).

• norm_first (bool) – whether to apply layer normalization before or after attention.

• max_block_size (int) – maximum context size.

• sa_kwargs (dict, optional) – keyword arguments for self-attention, by default None

• ra_kwargs (dict, optional) – keyword arguments for relational attention, by default None

• ra_type ('relational_attention', 'rca', or 'disrca', optional) – type of relational attention module (e.g., whether to use RCA for an ablation experiment), by default ‘relational_attention’

• share_attn_params (bool, optional) – whether to share attention parameters between self-attention and relational attention. If True, w{q,k} in sensory attention and w{q,k}_attn in relational attention are shared. number of heads in each must be the same. By default False

• symbol_retrieval ('symbolic_attention', 'position_relative', 'positional_symbols', optional) – type of symbol retrieval module to use. this is shared across layers, by default ‘symbolic_attention’

• pos_enc_type ('pos_emb' or 'RoPE', optional) – type of positional encoding to use, by default ‘pos_emb’

• bias (bool, optional) – whether to use bias in attention, by default True

forward(x, targets=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

estimate_mfu(fwdbwd_per_iter, dt)

estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS

get_num_params(non_embedding=True)

Return the number of parameters in the model. For non-embedding count (default), the position embeddings get subtracted.

generate(idx, max_new_tokens, temperature=1.0, top_k=None)

Generate max_new_tokens new tokens, conditioning on the input idx.

Parameters:

• idx (Tensor[int]) – tensor of shape (batch_size, seq_len) with input tokens.

• max_new_tokens (int) – number of new tokens to generate

• temperature (float, optional) – temperature parameter of softmax, by default 1.0

• top_k (int, optional) – top-k sampling parameter, by default None

Returns:

tensor of shape (batch_size, seq_len + max_new_tokens) with generated tokens.

Return type:

Tensor[int]

class dual_attention.language_models.TransformerLM(vocab_size: int, d_model: int, n_layers: int, n_heads: int, dff: int, dropout_rate: float, activation: str, norm_first: bool, max_block_size: int, norm_type: str = 'layernorm', bias: bool = True, pos_enc_type: str = 'pos_emb', use_flash_attention=True, block_kwargs: dict = None)

Bases: Module

Transformer Language Model

__init__(vocab_size: int, d_model: int, n_layers: int, n_heads: int, dff: int, dropout_rate: float, activation: str, norm_first: bool, max_block_size: int, norm_type: str = 'layernorm', bias: bool = True, pos_enc_type: str = 'pos_emb', use_flash_attention=True, block_kwargs: dict = None)

Transformer autoregressive language model.

given (x_1, …, x_T) causally predicts (y_1, …, y_T)

Parameters:

• vocab_size (int) – vocabulary size.

• d_model (int) – model dimension.

• n_layers (int) – number of layers.

• n_heads (int) – number of attention heads.

• dff (int) – size of intermediate layer in feedforward blocks.

• dropout_rate (float) – dropout rate.

• activation (str) – name of activation function (e.g., ‘relu’, ‘gelu’, or ‘swiglu’).

• norm_first (bool) – whether to apply layer normalization before or after attention.

• max_block_size (int) – maximum context size.

• bias (bool, optional) – whether to use bias in attention, by default True

• pos_enc_type ('pos_emb' or 'RoPE', optional) – type of positional encoding to use, by default ‘pos_emb’

forward(x, targets=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

estimate_mfu(fwdbwd_per_iter, dt)

estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS

get_num_params(non_embedding=True)

Return the number of parameters in the model. For non-embedding count (default), the position embeddings get subtracted.

generate(idx, max_new_tokens, temperature=1.0, top_k=None)

Generate max_new_tokens new tokens, conditioning on the input idx.

Parameters:

• idx (Tensor[int]) – tensor of shape (batch_size, seq_len) with input tokens.

• max_new_tokens (int) – number of new tokens to generate

• temperature (float, optional) – temperature parameter of softmax, by default 1.0

• top_k (int, optional) – top-k sampling parameter, by default None

Returns:

tensor of shape (batch_size, seq_len + max_new_tokens) with generated tokens.

Return type:

Tensor[int]

dual_attention.language_models.configure_optimizers(model, weight_decay, learning_rate, betas, device_type)

dual_attention.model_utils module

dual_attention.model_utils.get_activation_function(name)

gets activation function by its name.

dual_attention.positional_encoding module

class dual_attention.positional_encoding.SinusoidalPositionalEncoding(d_model: int, scale: bool = True, dropout: float = 0.1, max_len: int = 2048)

Bases: Module

__init__(d_model: int, scale: bool = True, dropout: float = 0.1, max_len: int = 2048)

module which adds a (non-trainable) sinusoidal positional encoding to the input tensor

Parameters:

• d_model (int) – model dimension.

• scale (bool, optional) – whether to scale added positional encodings to account for scaling in dot product attention, by default True

• dropout (float, optional) – dropout rate, by default 0.1

• max_len (int, optional) – maximum length to consider, by default 2048

forward(x: Tensor) → Tensor

Parameters:

x – Tensor, shape [seq_len, batch_size, embedding_dim]

class dual_attention.positional_encoding.LearnedPositionalEmbeddings(d_model: int, scale: bool = True, dropout: float = 0.1, max_len: int = 2048)

Bases: Module

__init__(d_model: int, scale: bool = True, dropout: float = 0.1, max_len: int = 2048)

module which adds a learnable positionall embedding to the input tensor.

Parameters:

• d_model (int) – model dimension.

• scale (bool, optional) – whether to scale added positional encodings to account for scaling in dot product attention, by default True

• dropout (float, optional) – dropout rate, by default 0.1

• max_len (int, optional) – maximum length to consider, by default 2048

forward(x: Tensor) → Tensor

Parameters:

x – Tensor, shape [seq_len, batch_size, embedding_dim]

class dual_attention.positional_encoding.RelativePositionalEncoding(dim: int, max_rel_pos: int)

Bases: Module

__init__(dim: int, max_rel_pos: int)

module which returns relative positional embeddings for a given pair of sequences.

I.e., returns tensor whose [i,j]-th entry is the embedding of the relative position “j-i”

Parameters:

• dim (int) – dimension of embeddings

• max_rel_pos (int) – maximum relative position in either direction (used for clipping)

forward(length_q, length_k=None)

Parameters:

• length_q (int) – length of query sequence

• length_k (_type_, optional) – length of key sequence, by default None

Returns:

tensor of shape [length_q, length_k, dim] where [i,j] is the embedding of the relative position “j-i”

Return type:

Tensor

dual_attention.relational_attention module

A module implementating relational attention.

This module also includes an implementation of relational cross-attention from our previous paper for ablation experiments.

class dual_attention.relational_attention.RelationalAttention(d_model: int, n_heads: int, n_relations: int = None, dropout: float = 0.0, key_dim: int = None, n_kv_heads: int = None, rel_activation: str = 'identity', rel_proj_dim: int = None, add_bias_kv: bool = False, add_bias_out: bool = False, total_n_heads: int = None, symmetric_rels: bool = False, use_relative_positional_symbols: bool = False)

Bases: Module

__init__(d_model: int, n_heads: int, n_relations: int = None, dropout: float = 0.0, key_dim: int = None, n_kv_heads: int = None, rel_activation: str = 'identity', rel_proj_dim: int = None, add_bias_kv: bool = False, add_bias_out: bool = False, total_n_heads: int = None, symmetric_rels: bool = False, use_relative_positional_symbols: bool = False)

An implementation of Relational Attention (RA).

Relational attention defines a differentiable information-retrieval operation where the information retrieved is the relations between objects. The “message” sent from one object to another is the relations between the sender and the receiver, tagged with a symbol identifying the sender. These messages are aggregated based on the receiver’s features via softmax attention scores.

The learnable parameters include a set of query/key projections which determine the attention scores, and hence the “selection criteria”, as well as a set of query/key projections for computing relations between objects. They also include per-head projections for the symbols and relations, as well as a final output projection.

This module supports symmetric relations, position-relative symbolic embeddings, multi-query attention/grouped query attention, and control over total number of heads (for use with “dual attention”).

Parameters:

• d_model (int) – model dimension

• n_heads (int) – number of attention heads (query heads if n_kv_heads is set)

• n_relations (int, optional) – number of relations. If None, n_relations = n_heads. By default None

• dropout (float, optional) – dropout rate. By default 0.0

• n_kv_heads (int, optional) – number of key/value heads. used to implement multi-query attention or grouped query attention. n_kv_heads=1 corresponds to MQA, n_kv_heads > 1 corresponsd to grouped query attention. n_kv_heads=n_heads is standard MHA. uses MHA when None. By default None

• rel_activation (str, optional) – name of activation function applied to relations. By default ‘identity’.

• rel_proj_dim (int, optional) – dimension of relation projections. If None, rel_proj_dim = head_dim * n_heads // n_relations. By default None.

• add_bias_kv (bool, optional) – whether to use bias in key/value projections, by default False

• add_bias_out (bool, optional) – whether to use bias in out projection, by default False

• total_n_heads (int, optional) – total number of heads in dual attention (if using dual attention). used to ensure that concat(A, E) is of dimension d_model after concatentation. hence, output dimension is (d_model // total_heads) * n_heads. if None, total_heads = n_heads and output dimension is d_model

forward(x: Tensor, symbols: Tensor, freqs_cos: Tensor = None, freqs_sin: Tensor = None, attn_mask: Tensor = None, is_causal: bool = False)

compute relational attention value.

if freqs_cos and freqs_sin are given, apply rotary positional embeddings.

if attn_mask is given, apply attention mask.

if is_causal is True, apply causal mask (attn_mask must be None).

if use_relative_positional_symbols is True, the symbols are treated as relative positional embeddings.

assumed to be of shape [len, len, dim] where len is the length of the sequence x.

Parameters:

• x (torch.Tensor) – input tensor of shape [bsz, len, d_model]

• symbols (torch.Tensor) – input tensor of shape [bsz, len, d_model] or [len, len, d_model] if use_relative_positional_symbols is True

• freqs_cos (torch.Tensor, optional) – cosine of frequencies for RoPE. RoPE is applied if given. By default None

• freqs_sin (torch.Tensor, optional) – cosine of frequencies for RoPE. RoPE is applied if given. By default None

• attn_mask (torch.Tensor, optional) – boolean attention mask of shape [len, len]. True at [i,j] indicates i is allowed to attend to j. By default None

• is_causal (bool, optional) – whether to apply a causal mask. If True, attn_mask must be None. By default False

Returns:

outputs [bsz, len, d_model], attention scores [bsz, n_heads, len, len], relations [bsz, len, len, n_relations]

Return type:

tuple[torch.Tensor]

class dual_attention.relational_attention.RelationalCrossAttention(d_model: int, n_heads: int, dropout: float, n_kv_heads: int = None, activation: str = 'softmax', add_bias_kv: bool = False, add_bias_out: bool = False, total_n_heads: int = None, use_relative_positional_symbols: bool = False)

Bases: Module

__init__(d_model: int, n_heads: int, dropout: float, n_kv_heads: int = None, activation: str = 'softmax', add_bias_kv: bool = False, add_bias_out: bool = False, total_n_heads: int = None, use_relative_positional_symbols: bool = False)

An implementation of Relational Cross Attention with some added customization.

Supports position-relative symbolic embeddings, multi-query attention/grouped query attention, and control over total number of heads (for use with “abstract attention”).

This corresponds to RCA as proposed by “Abstractors and relational cross-attention: An inductive bias for explicit relational reasoning in Transformers” Awni Altabaa, Taylor Webb, Jonathan Cohen, John Lafferty. ICLR (2024). https://arxiv.org/abs/2304.00195

Parameters:

• d_model (int) – model dimension

• n_heads (int) – number of heads (query heads if n_kv_heads is set)

• dropout (float) – dropout rate

• n_kv_heads (int, optional) – number of key/value heads. used to implement multi-query attention or grouped query attention. n_kv_heads=1 corresponds to MQA, n_kv_heads > 1 corresponsd to grouped query attention. n_kv_heads=n_heads is standard MHA. uses MHA when None. By default None

• activation (str, optional) – name of activation function applied to attention scores. If softmax, flash attention is used. Otherwise, attention is computed ‘manually’ with the chosen activation function. By default ‘softmax’.

• add_bias_kv (bool, optional) – whether to use bias in key/value projections, by default False

• add_bias_out (bool, optional) – whether to use bias in out projection, by default False

• total_n_heads (int, optional) – total number of heads in abstract attention (if using abstract attention). used to ensure that concat(A, E) is of dimension d_model after concatentation. hence, output dimension is (d_model // total_heads) * n_heads. if None, total_heads = n_heads and output dimension is d_model

• use_relative_pos_symbols (bool, optional) – whether to use relative positional symbols, by default False

forward(x: Tensor, symbols: Tensor, freqs_cos: Tensor = None, freqs_sin: Tensor = None, attn_mask: Tensor = None, is_causal: bool = False)

compute relational cross-attention attention with given input x and symbols.

if attn_mask is given, apply attention mask.

if is_causal is True, apply causal mask (attn_mask must be None).

if use_relative_pos_symbols is True, the symbols are treated as relative positional embeddings.

assumed to be of shape [len, len, dim] where len is the length of the sequence x.

Parameters:

• x (torch.Tensor) – input tensor of shape [bsz, len, d_model]

• symbols (torch.Tensor) – input tensor of shape [bsz, len, d_model] or [len, len, d_model] if use_relative_positional_symbols is True

• freqs_cos (torch.Tensor, optional) – cosine of frequencies for RoPE. RoPE is applied if given. By default None

• freqs_sin (torch.Tensor, optional) – cosine of frequencies for RoPE. RoPE is applied if given. By default None

• attn_mask (torch.Tensor, optional) – boolean attention mask of shape [len, len]. True at [i,j] indicates i is allowed to attend to j. By default None

• is_causal (bool, optional) – whether to apply a causal mask. If True, attn_mask must be None. By default False

Returns:

result of attention

Return type:

torch.Tensor

class dual_attention.relational_attention.DisentangledRelationalCrossAttention(d_model: int, n_heads: int, dropout: float, n_kv_heads: int = None, rel_activation: str = 'identity', add_bias_kv: bool = False, add_bias_out: bool = False, total_n_heads: int = None, use_relative_positional_symbols: bool = False)

Bases: Module

__init__(d_model: int, n_heads: int, dropout: float, n_kv_heads: int = None, rel_activation: str = 'identity', add_bias_kv: bool = False, add_bias_out: bool = False, total_n_heads: int = None, use_relative_positional_symbols: bool = False)

An implementation of Disentangled Relational Cross Attention with some added customization.

In Disentangled RCA, “attention” is separated from “relation”. Two sets of projections are learned, one for attention and one for relational representation. Attention scores determine which objects to attend to, and relation scores compute the relation between the objects, which is tied to the symbols to identify to the “sender”.

Supports position-relative symbolic embeddings, multi-query attention/grouped query attention, and control over total number of heads (for use with “abstract attention”).

Parameters:

• d_model (int) – model dimension

• n_heads (int) – number of heads (query heads if n_kv_heads is set)

• dropout (float) – dropout rate

• n_kv_heads (int, optional) – number of key/value heads. used to implement multi-query attention or grouped query attention. n_kv_heads=1 corresponds to MQA, n_kv_heads > 1 corresponsd to grouped query attention. n_kv_heads=n_heads is standard MHA. uses MHA when None. By default None

• rel_activation (str, optional) – name of activation function applied to relations. By default ‘identity’.

• add_bias_kv (bool, optional) – whether to use bias in key/value projections, by default False

• add_bias_out (bool, optional) – whether to use bias in out projection, by default False

• total_n_heads (int, optional) – total number of heads in abstract attention (if using abstract attention). used to ensure that concat(A, E) is of dimension d_model after concatentation. hence, output dimension is (d_model // total_heads) * n_heads. if None, total_heads = n_heads and output dimension is d_model

forward(x: Tensor, symbols: Tensor, freqs_cos: Tensor = None, freqs_sin: Tensor = None, attn_mask: Tensor = None, is_causal: bool = False)

compute attention with given query, key, value.

if freqs_cos and freqs_sin are given, apply rotary positional embeddings.

if attn_mask is given, apply attention mask.

if is_causal is True, apply causal mask (attn_mask must be None).

if use_relative_positional_symbols is True, the symbols are treated as relative positional embeddings.

assumed to be of shape [len, len, dim] where len is the length of the sequence x.

Parameters:

• x (torch.Tensor) – input tensor of shape [bsz, len, d_model]

• symbols (torch.Tensor) – input tensor of shape [bsz, len, d_model] or [len, len, d_model] if use_relative_positional_symbols is True

• freqs_cos (torch.Tensor, optional) – cosine of frequencies for RoPE. RoPE is applied if given. By default None

• freqs_sin (torch.Tensor, optional) – cosine of frequencies for RoPE. RoPE is applied if given. By default None

• attn_mask (torch.Tensor, optional) – boolean attention mask of shape [len, len]. True at [i,j] indicates i is allowed to attend to j. By default None

• is_causal (bool, optional) – whether to apply a causal mask. If True, attn_mask must be None. By default False

Returns:

result of attention

Return type:

torch.Tensor

dual_attention.seq2seq_models module

This module implements Encoder-Decoder Sequence-to-Sequence models (both Dual Attention Transformer and standard Transformer).

class dual_attention.seq2seq_models.Seq2SeqTransformer(input_spec: dict, output_spec: dict, d_model: int, out_dim: int, n_layers_enc: int, n_layers_dec: int, encoder_kwargs: dict, decoder_kwargs: dict, in_block_size: int, out_block_size: int, tie_weights: bool = True, loss_ignore_idx: int = -1)

Bases: Module

Transformer Language Model

__init__(input_spec: dict, output_spec: dict, d_model: int, out_dim: int, n_layers_enc: int, n_layers_dec: int, encoder_kwargs: dict, decoder_kwargs: dict, in_block_size: int, out_block_size: int, tie_weights: bool = True, loss_ignore_idx: int = -1)

Seq2Seq Encoder-Decoder Transformer.

Parameters:

• input_spec (dict) – description of input format. dictionary with key ‘type’ with values ‘token’ or ‘vector’. if ‘token’, must also have ‘vocab_size’. if ‘vector’, must also have ‘dim’.

• output_spec (dict) – description of output format. dictionary with key ‘type’ with values ‘token’ or ‘vector’. if ‘token’, must also have ‘vocab_size’. if ‘vector’, must also have ‘dim’.

• d_model (int) – model dimension.

• out_dim (int) – output dimension (e.g., output vocab size)

• n_layers_enc (int) – number of encoder layers.

• n_layers_dec (int) – number of decoder layers.

• encoder_kwargs (dict) – keyword arguments for encoder blocks.

• decoder_kwargs (dict) – keyword arguments for decoder blocks.

• in_block_size (int) – block size for input sequence.

• out_block_size (int) – block size for target sequence.

• tie_weights (bool, optional) – whether to tie weights between target embedder and final layer weights, by default True

• loss_ignore_idx (int, optional) – idx of class to ignore when computing loss, by default -1

forward(x, y, targets=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

get_num_params()

Return the number of parameters in the model.

estimate_mfu(fwdbwd_per_iter, dt)

estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS

class dual_attention.seq2seq_models.Seq2SeqDualAttnTransformer(input_spec: dict, output_spec: dict, symbol_retrieval: str, symbol_retrieval_kwargs: dict, d_model: int, out_dim: int, n_layers_enc: int, n_layers_dec: int, encoder_kwargs: dict, decoder_kwargs: dict, in_block_size: int, out_block_size: int, tie_weights: bool = True, loss_ignore_idx: int = -1)

Bases: Module

Dual Attention Transformer Seq2Seq Model

__init__(input_spec: dict, output_spec: dict, symbol_retrieval: str, symbol_retrieval_kwargs: dict, d_model: int, out_dim: int, n_layers_enc: int, n_layers_dec: int, encoder_kwargs: dict, decoder_kwargs: dict, in_block_size: int, out_block_size: int, tie_weights: bool = True, loss_ignore_idx: int = -1)

Seq2Seq Encoder-Decoder Dual Attention Transformer

Parameters:

• input_spec (dict) – description of input format. dictionary with key ‘type’ with values ‘token’ or ‘vector’. if ‘token’, must also have ‘vocab_size’. if ‘vector’, must also have ‘dim’

• output_spec (dict) – description of output format. dictionary with key ‘type’ with values ‘token’ or ‘vector’. if ‘token’, must also have ‘vocab_size’. if ‘vector’, must also have ‘dim’

• symbol_retrieval (str) – type of symbol retrieval mechanism. must be one of ‘symbolic_attention’, ‘rel_sym_attn’, ‘positional_symbols’, or ‘position_relative’

• symbol_retrieval_kwargs (dict) – keyword arguments for symbol retrieval mechanism

• d_model (int) – model dimension

• out_dim (int) – output dimension (e.g., output vocab size)

• n_layers_enc (int) – number of encoder layers

• n_layers_dec (int) – number of decoder layers

• encoder_kwargs (dict) – keyword arguments for encoder blocks

• decoder_kwargs (dict) – keyword arguments for decoder blocks

• in_block_size (int) – block size for input sequence

• out_block_size (int) – block size for target sequence

• tie_weights (bool, optional) – whether to tie weights between target embedder and final layer weights, by default True

• loss_ignore_idx (int, optional) – idx of class to ignore when computing loss, by default -1

forward(x, y, targets=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

get_num_params()

Return the number of parameters in the model.

estimate_mfu(fwdbwd_per_iter, dt)

estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS

dual_attention.seq2seq_models.configure_optimizers(model, weight_decay, learning_rate, betas, device_type)

dual_attention.symbol_retrieval module

This module implements various symbol assignment mechanisms (aka SymbolRetrievers) for Dual Attention Transformer models.

class dual_attention.symbol_retrieval.SymbolicAttention(d_model: int, n_heads: int, n_symbols: int, dropout: float = 0.0, scale: float = None, trainable_symbols: bool = True)

Bases: Module

__init__(d_model: int, n_heads: int, n_symbols: int, dropout: float = 0.0, scale: float = None, trainable_symbols: bool = True)

Symbolic Attention.

Learns a library of “symbols” and corresponding template features. For a given input, retrieves a symbol from the symbol library via attention.

Parameters:

• d_model (int) – model dimension. this is the dimension of the input and the dimension of the symbols and template features.

• n_heads (int) – number of heads in symbolic attention.

• n_symbols (int) – number of symbols in the symbol library.

• dropout (float, optional) – dropout probability, by default 0.0

• scale (float, optional) – scaling factor in scaled_dot_product_attention, by default None

• trainable_symbols (bool, optional) – whether to make the symbol library trainable, by default True

reset_parameters()

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dual_attention.symbol_retrieval.PositionalSymbolRetriever(symbol_dim, max_length, sinusoidal=False)

Bases: Module

__init__(symbol_dim, max_length, sinusoidal=False)

Postional Symbol Retriever.

Learns a library of “symbols”. Retrieves a symbol for each object based on its position.

Parameters:

• symbol_dim (int) – dimension of the symbols.

• max_symbols (int) – maximum number of symbols.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dual_attention.symbol_retrieval.PositionRelativeSymbolRetriever(symbol_dim, max_rel_pos)

Bases: Module

__init__(symbol_dim, max_rel_pos)

Position-Relative Symbol Retriever.

For i -> j, the symbol s_{ij} encodes the relative position j - i.

Parameters:

• symbol_dim (int) – dimension of the symbols.

• max_rel_pos (int) – maximum relative position encoded by symbols. Positions exceeding this will be truncated.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dual_attention.symbol_retrieval.RelationalSymbolicAttention(d_model: int, rel_n_heads: int, symbolic_attn_n_heads: int, n_symbols: int, nbhd_delta: int, causal_nbhd: bool = True, include_self: bool = False, normalize_rels: bool = True, dropout: float = 0.0, rel_scale: float = None, symbolic_attn_scale: float = None)

Bases: Module

__init__(d_model: int, rel_n_heads: int, symbolic_attn_n_heads: int, n_symbols: int, nbhd_delta: int, causal_nbhd: bool = True, include_self: bool = False, normalize_rels: bool = True, dropout: float = 0.0, rel_scale: float = None, symbolic_attn_scale: float = None)

Relational symbolic attention module.

Retrieves a symbol for each object in the input based on its relationship with its neighborhood. First, we compute a local relation vector for each object in the input. This local relation vector is then used to retrieve a symbol from the symbol library via symbolic attention.

Parameters:

• d_model (int) – Model dimension. this is the dimension of the input and the dimension of the symbols and template features.

• rel_n_heads (int) – Dimensionality of relations computed with neighborhood.

• symbolic_attn_n_heads (int) – Number of symbolic attention heads.

• n_symbols (int) – Number of symbols to learn in the symbol library.

• nbhd_delta (int) – The size of the neighborhood.

• causal_nbhd (bool, optional) – Whether to use causal neighborhood. if causal_nbhd is True, the neighborhood is [i-nbhd_delta, i]. if causal_nbhd is False, the neighborhood is [i-nbhd_delta, i+nbhd_delta]. Defaults to True.

• include_self (bool, optional) – Whether to include self in the neighborhood. E.g., if False and causal_nbhd, the neighborhood is [i-nbhd_delta, i-1]. If False and not causal_nbhd, the neighborhood is [i-nbhd_delta, i-1] U [i+1, i+nbhd_delta]. Defaults to False.

• normalize_rels (bool, optional) – Whether to normalize relations with softmax across neighborhood. Defaults to True.

• dropout (float, optional) – The dropout rate. Defaults to 0.0.

• rel_scale (float, optional) – The scaling factor when normalizing relations via softmax. If None, it is computed based on model_dim and rel_n_heads.

• symbolic_attn_scale (float, optional) – The scaling factor used in symbolic attention.

symbolic_attention

The symbolic attention module.

Type:

SymbolicAttention

q_proj

Linear layer for projecting the query.

Type:

nn.Linear

k_proj

Linear layer for projecting the key.

Type:

nn.Linear

model_dim_proj

Linear layer for projecting the neighborhood relation vector to model_dim.

Type:

nn.Linear

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

static compute_neighbor_mask(n, delta, include_self=True)

computes the neighborhood mask for a sequence of length n and neighborhood size delta

static compute_causal_neighbor_mask(n, delta, include_self=False)

computes the causal neighborhood mask for a sequence of length n and neighborhood size delta

dual_attention.transformer_blocks module

This module implements Encoder and Decoder blocks for standard Transformer architectures.

Each block consists of: Self-attention, (Cross-attention for DecoderBlock), Feed-forward block, LayerNorms/Residuals.

class dual_attention.transformer_blocks.EncoderBlock(d_model: int, n_heads: int, dff: int, dropout_rate: float, activation: str, norm_first: bool, norm_type: str = 'layernorm', bias: bool = True, causal: bool = False, attn_kwargs: dict = None)

Bases: Module

__init__(d_model: int, n_heads: int, dff: int, dropout_rate: float, activation: str, norm_first: bool, norm_type: str = 'layernorm', bias: bool = True, causal: bool = False, attn_kwargs: dict = None)

A Transformer Encoder Block.

Consists of Self-attention, Feed-forward block and LayerNorms/Residuals.

Parameters:

• d_model (int) – model dimension.

• n_heads (int) – number of self-attention heads.

• dff (int) – intermediate dimension of feed-forward block.

• dropout_rate (float) – dropout rate.

• activation (str) – name of activation function to use in feed-forward block.

• norm_first (bool) – whether to apply layer normalization before or after attention.

• norm_type (str, optional) – type of normalization to use. ‘layernorm’ or ‘rmsnorm’. Default is ‘layernorm’.

• bias (bool, optional) – whether to use bias in multi-head attention, by default True

• resgate_kwargs (dict, optional) – keyword arguments for ResidualGate, by default None

• causal (bool, optional) – whether self-attention should be causal, by default False

forward(x, freqs_cos=None, freqs_sin=None, need_weights=False)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dual_attention.transformer_blocks.DecoderBlock(d_model: int, n_heads: int, n_heads_cross: int, dff: int, dropout_rate: float, activation: str, norm_first: bool, norm_type: str = 'layernorm', bias: bool = True, causal: bool = False)

Bases: Module

__init__(d_model: int, n_heads: int, n_heads_cross: int, dff: int, dropout_rate: float, activation: str, norm_first: bool, norm_type: str = 'layernorm', bias: bool = True, causal: bool = False)

A Transformer Decoder Block.

Consists of Self-attention, Cross-attention, Feed-forward block and LayerNorms/Residuals.

Parameters:

• d_model (int) – model dimension.

• n_heads (int) – number of self-attention heads.

• n_heads_cross (int) – number of cross-attention heads.

• dff (int) – intermediate dimension of feed-forward block.

• dropout_rate (float) – dropout rate.

• activation (str) – name of activation function to use in feed-forward block.

• norm_first (bool) – whether to apply layer normalization before or after attention.

• norm_type (str, optional) – type of normalization to use. ‘layernorm’ or ‘rmsnorm’. Default is ‘layernorm’.

• bias (bool, optional) – whether to use bias in multi-head attention, by default True

• causal (bool, optional) – whether self-attention should be causal, by default False

forward(x, context)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dual_attention.transformer_blocks.FeedForwardBlock(embed_dim: int, dff: int = None, activation: str = 'relu', use_bias: bool = False)

Bases: Module

__init__(embed_dim: int, dff: int = None, activation: str = 'relu', use_bias: bool = False)

Feed-forward block.

A 2-layer neural network with activation function in between.

Parameters:

• embed_dim (int) – embedding dimension of input.

• dff (int, optional) – size of intermediate layer. if None, 4 * embed_dim.

• activation (str, optional) – name of activation function, by default ‘relu’

• use_bias (bool, optional) – whether to use bias in linear layers, by default False

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dual_attention.transformer_blocks.RMSNorm(dim: int, eps: float = 1e-05)

Bases: Module

__init__(dim: int, eps: float = 1e-05)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

dual_attention.transformer_blocks.create_norm(d_model, norm_type)

dual_attention.vision_models module

Implementation of Vision Transformer (ViT) and Vision Dual Attention Transformer (ViDAT)

class dual_attention.vision_models.VisionTransformer(image_shape: Tuple[int], patch_size: Tuple[int], num_classes: int, d_model: int, n_layers: int, n_heads: int, dff: int, dropout_rate: float, activation: str, norm_first: bool, norm_type: str = 'layernorm', bias: bool = True, pool: str = 'cls')

Bases: Module

Vision Transformer

__init__(image_shape: Tuple[int], patch_size: Tuple[int], num_classes: int, d_model: int, n_layers: int, n_heads: int, dff: int, dropout_rate: float, activation: str, norm_first: bool, norm_type: str = 'layernorm', bias: bool = True, pool: str = 'cls')

Vision Transformer.

Parameters:

• image_shape (Tuple[int]) – shape of image (channels, width, height)

• patch_size (Tuple[int]) – size of patch (width, height)

• num_classes (int) – number of classes

• d_model (int) – model dimension

• n_layers (int) – number of layers

• n_heads (int) – number of attention heads

• dff (int) – feedforward dimension

• dropout_rate (float) – dropout rate

• activation (str) – name of activation function in feedforward blocks

• norm_first (bool) – whether to apply normalization before or after attention. norm_first=True means pre-norm otherwise post-norm.

• norm_type ('layernorm' or 'rmsnorm', optional) – type of normalization to use, by default ‘layernorm’

• bias (bool, optional) – whether to use a bias in the encoder blocks, by default True

• pool ('cls' or 'mean', optional) – type of pooling to use before final class prediction. ‘cks’ corresponds to using a class token while ‘mean’ corresponds to mean pooling, by default ‘cls’

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dual_attention.vision_models.VisionDualAttnTransformer(image_shape: Tuple[int], patch_size: Tuple[int], num_classes: int, d_model: int, n_layers: int, n_heads_sa: int, n_heads_ra: int, dff: int, dropout_rate: float, activation: str, norm_first: bool, symbol_retrieval: str, symbol_retrieval_kwargs: dict, ra_type: str = 'relational_attention', ra_kwargs: dict = None, sa_kwargs: dict = None, norm_type: str = 'layernorm', bias: bool = True, pool: str = 'cls')

Bases: Module

Vision Dual Attention Transformer

__init__(image_shape: Tuple[int], patch_size: Tuple[int], num_classes: int, d_model: int, n_layers: int, n_heads_sa: int, n_heads_ra: int, dff: int, dropout_rate: float, activation: str, norm_first: bool, symbol_retrieval: str, symbol_retrieval_kwargs: dict, ra_type: str = 'relational_attention', ra_kwargs: dict = None, sa_kwargs: dict = None, norm_type: str = 'layernorm', bias: bool = True, pool: str = 'cls')

Vision Transformer.

Parameters:

• image_shape (Tuple[int]) – shape of image (channels, width, height)

• patch_size (Tuple[int]) – size of patch (width, height)

• num_classes (int) – number of classes

• d_model (int) – model dimension

• n_layers (int) – number of layers

• n_heads_sa (int) – number of self-attention heads

• n_heads_ra (int) – number of relational attention heads

• dff (int) – feedforward dimension

• dropout_rate (float) – dropout rate

• activation (str) – name of activation function in feedforward blocks

• norm_first (bool) – whether to apply normalization before or after attention. norm_first=True means pre-norm otherwise post-norm.

• symbol_retrieval (str) – type of symbol retrieval mechanism to use, one of ‘symbolic_attention’, ‘rel_sym_attn’, ‘positional_symbols’, ‘position_relative’

• symbol_retrieval_kwargs (dict) – keyword arguments for symbol retrieval mechanism

• ra_type ('relational_attention', 'rca', or 'disrca', optional) – type of relational attention module (e.g., whether to use RCA for an ablation experiment), by default ‘relational_attention’

• ra_kwargs (dict, optional) – relational attention kwargs, by default None

• sa_kwargs (dict, optional) – self-attention kwargs, by default None

• norm_type ('layernorm' or 'rmsnorm', optional) – type of normalization to use, by default ‘layernorm’

• bias (bool, optional) – whether to use a bias in the encoder blocks, by default True

• pool ('cls' or 'mean', optional) – type of pooling to use before final class prediction. ‘cks’ corresponds to using a class token while ‘mean’ corresponds to mean pooling, by default ‘cls’

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

dual_attention.vision_models.configure_optimizers(model, weight_decay, learning_rate, betas, device_type)

dual_attention.model_analysis.datlm_utils module

This module implements forward calls for the Dual Attention Transformer (DAT) model that return intermediate results for visualization purposes.

dual_attention.model_analysis.datlm_utils.symbolic_attn_forward_get_weights(mod, x)

a variant of the forward call for symbolic attention that returns the attention weights

dual_attention.model_analysis.datlm_utils.block_forward_get_weights(mod, x, symbols, freqs_cos=None, freqs_sin=None)

a variant of the forward call for a block that returns the attention weights

dual_attention.model_analysis.datlm_utils.datlm_forward_w_intermediate_results(model, x)

a variant of the forward call for the Dual Attention Transformer LM that returns intermediate results

dual_attention.model_analysis.lm_inference_app module

dual_attention.model_analysis.lm_visualization_app module