Induced Self-Attention Block for efficient O(n) attention

This module implements a bottleneck attention mechanism using a small set of learned inducing points that mediate interactions between input elements. The complexity is reduced from O(n²) to O(n) by:

1. Projecting inputs onto inducing points (size m << n)

2. Propagating information through these inducing points

3. Projecting back to the original sequence

Usage

induced_self_attention_block(
  d_model,
  nhead,
  dim_feedforward,
  num_inds,
  dropout = 0,
  activation = "gelu",
  norm_first = TRUE,
  bias_free_ln = FALSE,
  ssmax = FALSE,
  skip_value = -100
)

Arguments

d_model

Integer. Model dimension.

nhead

Integer. Number of attention heads.

dim_feedforward

Integer. Dimension of the feedforward network.

num_inds

Integer. Number of inducing points (controls capacity vs. efficiency).

dropout

Numeric. Dropout probability (default: 0.0).

activation

Character or function. Activation function for feedforward (default: "gelu").

norm_first

Logical. If TRUE, uses pre-norm architecture (default: TRUE).

bias_free_ln

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

ssmax

Logical or character. Type of scalable softmax to use in attention. Note that only the first attention layer uses SSMax (default: FALSE).

skip_value

Numeric. Value used to mark inputs that should be skipped (default: -100.0).

Value

A nn_module for induced self-attention.

References

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