Stack of induced self-attention blocks.

A set transformer uses induced self-attention mechanism to efficiently process variable-sized sets while maintaining permutation invariance.

Usage

set_transformer(
  num_blocks,
  d_model,
  nhead,
  dim_feedforward,
  num_inds = 16L,
  dropout = 0,
  activation = "gelu",
  norm_first = TRUE,
  bias_free_ln = FALSE,
  ssmax = FALSE,
  recompute = FALSE
)

Arguments

num_blocks

Integer. Number of induced self-attention blocks in the stack.

d_model

Integer. Model dimension.

nhead

Integer. Number of attention heads and should be a divisor of d_model.

dim_feedforward

Integer. Dimension of the feedforward network in each block.

num_inds

Integer, default 16L. Number of inducing points used in self-attention blocks.

dropout

Float, default 0. Dropout probability.

activation

Character string or unary function, default "gelu". The activation function used in the feedforward network.

norm_first

Logical, default TRUE. If TRUE, uses pre-norm architecture (LayerNorm before attention and feedforward).

bias_free_ln

Logical, default FALSE. If TRUE, removes bias from all LayerNorm layers.

ssmax

Logical or character string, default FALSE. Type of scalable softmax to use in attention. Note that only the first attention layer of the induced self-attention blocks uses SSMax. If TRUE, equivalent to "qassmax-mlp-elementwise". If FALSE, equivalent to "none". If a string, uses the specified scalable softmax type. Options include: "none", "ssmax", "ssmax-mlp", "ssmax-mlp-elementwise", "qassmax-mlp", "qassmax-mlp-elementwise".

recompute

Logical, default FALSE. If TRUE, uses gradient checkpointing to save memory at the cost of additional computation.

Value

An nn_module instance of class SetTransformer.

Methods

Usage

st <- SetTransformer(num_blocks = 3L, d_model = 128L, nhead = 8L,
                     dim_feedforward = 256L, num_inds = 128L)
st$forward(src, train_size = NULL)
st$forward_with_cache(src, col_cache, train_size, use_cache, store_cache)

References

Lee et al. "Set Transformer: A Framework for Attention-based Permutation-Invariant Neural Networks", ICML 2019.

See also

induced_self_attention_block, KVCache, KVCacheEntry