Row-Wise Feature Interaction Module

Captures interactions between features within each row using a transformer encoder with rotary positional encoding (RoPE). Prepends learnable CLS tokens to feature embeddings and uses only those tokens to aggregate per-row representations.

Usage

RowInteraction(
  embed_dim,
  num_blocks,
  nhead,
  dim_feedforward,
  num_cls = 4L,
  rope_base = 1e+05,
  rope_interleaved = TRUE,
  dropout = 0,
  activation = "gelu",
  norm_first = TRUE,
  bias_free_ln = FALSE,
  recompute = FALSE
)

Arguments

embed_dim

Integer. Embedding dimension.

num_blocks

Integer. Number of transformer blocks in the encoder.

nhead

Integer. Number of attention heads.

dim_feedforward

Integer. Dimension of the feedforward network.

num_cls

Integer. Number of learnable CLS tokens prepended to feature embeddings. Their outputs are concatenated to form the row representation (default: 4L).

rope_base

Numeric. Base scaling factor for rotary position encoding (default: 100000).

rope_interleaved

Logical. If TRUE, uses interleaved rotation where dimension pairs are (0,1), (2,3), etc. If FALSE, splits the embedding into first and second halves (default: TRUE).

dropout

Numeric. Dropout probability in the encoder (default: 0.0).

activation

Character or function. Activation for the feedforward network: "relu", "gelu", or a unary callable (default: "gelu").

norm_first

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

bias_free_ln

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

recompute

Logical. If TRUE, uses gradient checkpointing to save memory at the cost of extra computation (default: FALSE).

Value

An nn_module instance of class RowInteraction.

Forward signature

row_int <- RowInteraction(embed_dim = 128L, num_blocks = 3L,
                          nhead = 8L, dim_feedforward = 256L)
# Training
repr <- row_int(embeddings, d = d)
# Inference
repr <- row_int(embeddings)

See also

Encoder, RotaryEmbedding