Multi-Head Attention

Multi-head attention as described in Attention Is All You Need.

This is a thin wrapper around torch::nn_multihead_attention() that makes it usable as a building block of a Graph of tensor operations, where both self-attention and cross-attention can be expressed, see section Input and Output Channels.

Tensor Layout

All inputs and outputs are (batch, sequence, feature), i.e. the batch_first layout of torch::nn_multihead_attention(), which is fixed and not a hyperparameter. torch defaults to (sequence, batch, feature), but the first dimension of every shape has to be the batch dimension here.

nn_module

Calls torch::nn_multihead_attention() when trained, where the parameters embed_dim, kdim and vdim are inferred as the last dimension of the query, key and value tensors respectively, and batch_first is always TRUE, see section Tensor Layout.

Parameters

• num_heads :: integer(1)
  Number of parallel attention heads. The embedding dimension must be divisible by num_heads.

• dropout :: numeric(1)
  Dropout probability on the attention weights. Default is 0.

• bias :: logical(1)
  Whether to add a bias to the input and output projections. Default is TRUE.

• add_bias_kv :: logical(1)
  Whether to add a bias to the key and value sequences at dimension 1. Default is FALSE.

• add_zero_attn :: logical(1)
  Whether to add a new batch of zeros to the key and value sequences at dimension 1. Default is FALSE.

• avg_weights :: logical(1)
  Whether the returned attention weights are averaged over the attention heads. Default is TRUE. Only has an effect when "weights" is among $outputs.

Note that embed_dim, kdim and vdim are not parameters, as they are inferred from the shapes of the input tensors, and that batch_first is not a parameter either, as it is fixed to TRUE, see section Tensor Layout.

Input and Output Channels

The number of input channels is determined by the construction argument mode:

• mode = "self" (default): one input channel "input", which is used as query, key and value, i.e. the PipeOp performs self-attention.

• mode = "cross": input channels "query" and "key_value", i.e. the PipeOp performs cross-attention, where the second input is used as both key and value.

• mode = "general": input channels "query", "key" and "value", i.e. the PipeOp performs cross-attention with separate key and value inputs.

The module has two outputs, "output", containing the attention output, and "weights", containing the attention weights, of which only "output" is an output channel by default. Set $outputs to also (or only) get the weights, e.g. outputs = c("output", "weights"), which are only computed when they are among $outputs.

For an explanation see PipeOpTorch.

State

The state is the value calculated by the public method $shapes_out().

References

Vaswani A, Shazeer N, Parmar N, Uszkoreit J, Jones L, Gomez A, Kaiser Ł, Polosukhin I (2017). “Attention is all you need.” Advances in neural information processing systems, 30.

Super classes

mlr3pipelines::PipeOp -> PipeOpTorch -> PipeOpTorchMultiheadAttention

Methods

Public methods

• PipeOpTorchMultiheadAttention$new()

• PipeOpTorchMultiheadAttention$clone()

Inherited methods

• mlr3pipelines::PipeOp$help()
• mlr3pipelines::PipeOp$predict()
• mlr3pipelines::PipeOp$print()
• mlr3pipelines::PipeOp$train()
• PipeOpTorch$shapes_out()

---

PipeOpTorchMultiheadAttention$new()

Creates a new instance of this R6 class.

Usage

PipeOpTorchMultiheadAttention$new(
  id = "nn_multihead_attention",
  mode = "self",
  param_vals = list()
)

Arguments

id

(character(1))
Identifier of the resulting object.

mode

(character(1))
The attention mode, which determines the input channels. One of "self", "cross" or "general". This is a construction argument (and not a hyperparameter), because it determines the structure of the Graph. The default is "self", which means that the PipeOp performs self-attention. See section Input and Output Channels for more information.

param_vals

(list())
List of hyperparameter settings, overwriting the hyperparameter settings that would otherwise be set during construction.

---

PipeOpTorchMultiheadAttention$clone()

The objects of this class are cloneable with this method.

Usage

PipeOpTorchMultiheadAttention$clone(deep = FALSE)

Arguments

deep

Whether to make a deep clone.

Examples

# Construct the PipeOp
pipeop = nn("multihead_attention", num_heads = 4)
pipeop
#> 
#> ── PipeOp <multihead_attention>: not trained ───────────────────────────────────
#> Values: num_heads=4
#> 
#> ── Input channels: 
#>    name           train predict
#>  <char>          <char>  <char>
#>   input ModelDescriptor    Task
#> 
#> ── Output channels: 
#>    name           train predict
#>  <char>          <char>  <char>
#>  output ModelDescriptor    Task
# The available parameters
pipeop$param_set
#> <ParamSet(6)>
#>               id    class lower upper nlevels        default  value
#>           <char>   <char> <num> <num>   <num>         <list> <list>
#> 1:     num_heads ParamInt     1   Inf     Inf <NoDefault[0]>      4
#> 2:       dropout ParamDbl     0     1     Inf              0 [NULL]
#> 3:          bias ParamLgl    NA    NA       2           TRUE [NULL]
#> 4:   add_bias_kv ParamLgl    NA    NA       2          FALSE [NULL]
#> 5: add_zero_attn ParamLgl    NA    NA       2          FALSE [NULL]
#> 6:   avg_weights ParamLgl    NA    NA       2           TRUE [NULL]