Distribution-aware Column-wise Embedding Module

This module maps scalar cells to high-dimensional embeddings using a shared Set Transformer. It supports feature grouping, target-aware encoding, and affine transformation of features.

Usage

ColEmbedding(
  embed_dim,
  num_blocks,
  nhead,
  dim_feedforward,
  num_inds,
  dropout = 0,
  activation = "gelu",
  norm_first = TRUE,
  bias_free_ln = FALSE,
  affine = TRUE,
  feature_group = FALSE,
  feature_group_size = 3L,
  target_aware = FALSE,
  max_classes = 10L,
  reserve_cls_tokens = 4L,
  ssmax = FALSE,
  mixed_radix_ensemble = TRUE,
  recompute = FALSE
)

Arguments

embed_dim

Integer. Embedding dimension.

num_blocks

Integer. Number of induced self-attention blocks.

nhead

Integer. Number of attention heads.

dim_feedforward

Integer. Feedforward network dimension.

num_inds

Integer. Number of inducing points.

dropout

Numeric. Dropout probability (default: 0.0).

activation

Character. Activation function (default: "gelu").

norm_first

Logical. Pre-norm architecture (default: TRUE).

bias_free_ln

Logical. Remove LayerNorm bias (default: FALSE).

affine

Logical. Compute features * W + b (default: TRUE).

feature_group

Logical or Character. Grouping mode:

• FALSE: No grouping.

• TRUE or "same": Circular permutation.

• "valid": Padding and reshaping.

feature_group_size

Integer. Size of feature groups (default: 3L).

target_aware

Logical. Use target info (default: FALSE).

max_classes

Integer. Number of classes (default: 10L).

reserve_cls_tokens

Integer. CLS token slots (default: 4L).

ssmax

Logical or Character. Scalable softmax type.

mixed_radix_ensemble

Logical. Mixed-radix ensembling (default: TRUE).

recompute

Logical. Gradient checkpointing (default: FALSE).

Value

An nn_module object.