Layer Classes

Layer Classes§

The Dual Attention Transformer (DAT) architecture is an extension of the Transformer framework which introduces several new computational mechanisms. The dual-attention package implements these new mechanisms as PyTorch nn.Module classes. We recommend reading the paper first to better understand the package and documentation. In this section, we highlight some key classes and their usage.

Relational Attention§

Relational attention is an attentional mechanism proposed as part of the DAT architecture which routes relational information between elements of the input. It takes the following form:

\[\begin{split}\begin{align*} \text{RelationalAttention}(x, (y_1, \ldots, y_n)) &= \sum_{i} \alpha_{i}(x, \boldsymbol{y}) (r(x, y_i) W_r + s_i W_s), \\ \alpha(x, \boldsymbol{y}) &= \text{softmax}([\langle \phi_q^{\mathrm{attn}}(x), \phi_k^{\mathrm{attn}}(y_i) \rangle]_{i=1}^{n}), \\ r(x, y) &= (\langle \phi_{q,\ell}^{\mathrm{rel}}(x), \phi_{k, \ell}^{\mathrm{rel}}(y))_{\ell \in [d_r]}, \\ (s_1, \ldots, s_n) &= \text{SymbolRetriever}(\boldsymbol{y}; S_{\mathrm{lib}}), \end{align*}\end{split}\]

The RelationalAttention class in the dual_attention.relational_attention module implements the relational attention mechanism. It can be instantiated and used as a PyTorch layer as follows:

from dual_attention.relational_attention import RelationalAttention

relational_attn = RelationalAttention(
    d_model=64,                # model dimension
    n_heads=4,                 # number of heads
    n_relations=16,            # number of relations
    dropout=0.0,               # dropout rate
    key_dim=None,              # key dimension
    n_kv_heads=None,           # number of key-value heads (e.g., for MQA or GQA)
    rel_activation='identity', # activation function applied to relations
    rel_proj_dim=None,         # relation projection dimension (analogous to key_dim)
    add_bias_kv=False,         # whether to add bias to key-value projection
    add_bias_out=False,        # whether to add bias to output projection
    total_n_heads=None,        # total number of heads of DualAttn (if used within DualAttn)
    symmetric_rels=False,      # whether relations should be constrained to be symmetric
    use_relative_positional_symbols=False # whether to use relative positional symbols
    )
x = torch.randn(1, 10, 64)
s = torch.randn(1, 10, 64)
out, *_ = relational_attn(x, s)

print(f'x: {x.shape}, s: {s.shape}, out: {out.shape}')
# x: torch.Size([1, 10, 64]), s: torch.Size([1, 10, 64]), out: torch.Size([1, 10, 64])

Symbol Assignment Mechanisms§

The dual_attention.symbol_retrieval module implements several symbol assignment mechanisms which can be used to assign symbols to elements of the input in relational attention. The following symbol assignment mechanisms are implemented:

  • SymbolicAttention: Assigns symbols based on an attention operation to a library of symbols, learned as parameters of the model. This corresponds to computing a type of soft “equivalence class” over the input elements.

  • PositionalSymbolRetriever: Elements of the input are assigned a symbol based on their position in the input sequence (essentially a positional embedding).

  • PositionRelativeSymbolRetriever: Elements of the input are assigned different symbols depending on the “receiver”, encoding the relative of the sender with respect to the receiver. To use this, the use_relative_positional_symbols argument of the RelationalAttention module must be set to True.

Dual Attention§

Dual attention is a variant of multi-head attention in which there exists two sets of attention heads of different types: standard self-attention and relational attention. The two sets of attention heads are computed in parallel and concatenated to form the final output. This equips the model with two versatile computational mechanisms for processing and attending over both sensory and relational information.

The DualAttention class in the dual_attention.dual_attention module implements the dual attention mechanism. It can be instantiated and used as a PyTorch layer as follows:

from dual_attention.dual_attention import DualAttention

dual_attention = DualAttention(
    d_model=64,                # model dimension
    n_heads_sa=4,              # number of self-attention heads
    n_heads_ra=4,              # number of relational attention heads
    dropout=0.,                # dropout rate
    sa_kwargs=None,            # extra kwargs for self-attention
    ra_kwargs=None,            # extra kwargs for relational attention
    share_attn_params=False,   # whether to share attention parameters between self-attention and relational attention
    ra_type='relational_attention', # type of relational attention to use; used in ablation experiments in paper
    )

x = torch.randn(1, 10, 64)
s = torch.randn(1, 10, 64)
out, *_ = dual_attention(x, s)

print(f'x: {x.shape}, s: {s.shape}, out: {out.shape}')
# x: torch.Size([1, 10, 64]), s: torch.Size([1, 10, 64]), out: torch.Size([1, 10, 64])

Dual Attention Transformer Blocks§

The blocks in a Dual Attention Transformer architecture mirror those of a standard Transformer, but with dual attention replacing standard multi-head self-attention. They consist of dual-attention followed by a feedforward block, with layer normalization and residual connections around each sub-block.

The DualAttentionEncoderBlock and DualAttentDecoderBlock classes in the dual_attention.dual_attn_blocks module implement Dual Attention Transformer blocks. They can be instantiated and used as PyTorch layers as follows:

from dual_attention.dual_attn_blocks import DualAttnEncoderBlock

dat_block = DualAttnEncoderBlock(
    d_model=64,                # model dimension
    n_heads_sa=4,              # number of self-attention heads
    n_heads_ra=4,              # number of relational attention heads
    dff=256,                   # feedforward intermediate dimension
    activation='gelu',         # activation function of feedforward block
    dropout_rate=0.,           # dropout rate
    norm_first=True,           # whether to use pre-norm or post-norm
    norm_type='layernorm',     # type of normalization to use (e.g., 'layernorm' or 'rmsnorm')
    sa_kwargs=None,            # extra kwargs for self-attention
    ra_kwargs=None,            # extra kwargs for relational attention
    share_attn_params=False,   # whether to share attention parameters between self-attention and relational attention
    ra_type='relational_attention', # type of relational attention to use; used in ablation experiments in paper
    bias=True,                 # whether to use bias
    causal=False               # whether to use causal masking (e.g., for use in language model)
    )

x = torch.randn(1, 10, 64)
s = torch.randn(1, 10, 64)
out, *_ = dat_block(x, s)

print(f'x: {x.shape}, s: {s.shape}, out: {out.shape}')
# x: torch.Size([1, 10, 64]), s: torch.Size([1, 10, 64]), out: torch.Size([1, 10, 64])
Contents