Dataset-wise In-Context Learning Module

Implements in-context learning for tabular data using transformer architecture. Supports both classification (with hierarchical classification for many classes) and regression tasks. Includes KV caching for efficient inference.

Usage

ICLearning(
  max_classes,
  out_dim,
  d_model,
  num_blocks,
  nhead,
  dim_feedforward,
  dropout = 0,
  activation = "gelu",
  norm_first = TRUE,
  bias_free_ln = FALSE,
  ssmax = FALSE,
  recompute = FALSE
)

Arguments

max_classes

Integer. Determines task type and output behavior:

• If max_classes = 0: Regression using quantile prediction

• If max_classes > 0: Classification. Specifies native class support. If dataset classes exceed this, hierarchical classification is used.

out_dim

Integer. Output dimension of the model.

d_model

Integer. Model dimension for the transformer.

num_blocks

Integer. Number of transformer blocks.

nhead

Integer. Number of attention heads.

dim_feedforward

Integer. Feedforward network dimension.

dropout

Numeric. Dropout probability (default: 0.0).

activation

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

norm_first

Logical. Use pre-norm architecture (default: TRUE).

bias_free_ln

Logical. Remove bias from LayerNorm (default: FALSE).

ssmax

Logical or character. Scalable softmax type (default: FALSE).

recompute

Logical. Use gradient checkpointing (default: FALSE).

Value

An nn_module object.

Methods

forward(R, y_train, ...)

Main forward pass

prepare_repr_cache(R, y_train)

Prepare representation cache

forward_with_repr_cache(...)

Forward with representation cache

forward_with_cache(...)

Forward with KV cache

Examples

if (FALSE) { # \dontrun{
icl <- ICLearning(
  max_classes = 10L,
  out_dim = 10L,
  d_model = 128L,
  num_blocks = 6L,
  nhead = 8L,
  dim_feedforward = 512L
)
} # }