Skip to contents

TabICL: A Tabular In-Context Learning Foundation Model

Source: R/tab_icl2_full_model.R
TabICLv2.Rd

TabICL is a transformer-based architecture for in-context learning on tabular data. This class is the underlying raw torch module.

Usage

TabICLv2(
  config,
  max_classes = 10L,
  col_affine = FALSE,
  col_feature_group = "same",
  col_feature_group_size = 3L,
  col_target_aware = TRUE,
  col_ssmax = "qassmax-mlp-elementwise",
  row_rope_base = 1e+05,
  row_rope_interleaved = FALSE,
  icl_ssmax = "qassmax-mlp-elementwise",
  ff_factor = 2L,
  dropout = 0,
  activation = "gelu",
  norm_first = TRUE,
  bias_free_ln = FALSE,
  recompute = FALSE
)

Arguments

config

A tab_icl2_config object providing

  • embed_dim Integer, default 128L. Model dimension.

  • num_quantiles Integer, default 999L. Number of quantiles for regression.

  • col_n_block Integer, default 3L.

  • col_n_head Integer, default 8L.

  • col_n_cls Integer, default 128L.

  • col_feature_group_size Integer, default 3L.

  • row_n_block Integer, default 3L.

  • row_n_head Integer, default 8L.

  • row_n_cls Integer, default 4L.

  • icl_n_block Integer, default 12L.

  • icl_n_head Integer, default 8L.

max_classes

Integer, default 10L. 0 for regression.

col_affine

Logical, default FALSE.

col_feature_group

Character, default "same".

col_target_aware

Logical, default TRUE.

col_ssmax

Character, default "qassmax-mlp-elementwise".

row_rope_base

Float, default 100000.

row_rope_interleaved

Logical, default FALSE.

icl_ssmax

Character, default "qassmax-mlp-elementwise".

ff_factor

Integer, default 2L.

dropout

Float, default 0.

activation

Character, default "gelu".

norm_first

Logical, default TRUE.

bias_free_ln

Logical, default FALSE.

recompute

Logical, default FALSE.

Value

An nn_module instance of class TabICL.

Details

TabICL: A Tabular In-Context Learning Foundation Model

TabICL is a transformer-based architecture for in-context learning on tabular data to make predictions without fine-tuning. It processes tabular data through three sequential stages:

  1. Column-wise embedding creates distribution-aware embeddings.

  2. Row-wise interaction captures interactions between features within each row.

  3. Dataset-wise in-context learning to learn patterns from labeled examples and make predictions.

This class is the underlying raw torch module for TabICL. It is not intended to be used directly. Instead, use the classes from the top-level tabicl package such as TabICLClassifier or TabICLRegressor that wrap this class to include the necessary preprocessing of input features and postprocessing of predictions.

Methods

Usage


model <- TabICL$new(max_classes = 10L)
model$forward(X, y_train)
model$predict_stats(X, y_train, output_type = "mean")
model$forward_with_cache(X_train, y_train, store_cache = TRUE)
model$predict_stats_with_cache(X_test = X_test, use_cache = TRUE)
model$has_cache()
model$clear_cache()

forward(X, y_train, d, embed_with_test, feature_shuffles, return_logits, softmax_temperature, inference_config)

Column-wise embedding -> row-wise interaction -> dataset-wise in-context learning. Dispatches to train_forward() in training mode and inference_forward() in evaluation mode.

  • X: Input tensor of shape (B, TT, H) where B is the number of tables, TT is the number of samples (rows), and H is the number of features (columns). The first train_size positions contain training samples, and the remaining positions contain test samples.

  • y_train: Training labels of shape (B, train_size).

  • d: Optional tensor. The number of features per dataset. Used only in training mode.

  • embed_with_test: Logical, default FALSE. If TRUE, allow training samples to attend to test samples during embedding.

  • feature_shuffles: Optional list of integer vectors. Feature shuffle patterns for each table in the batch. Used only in inference mode.

  • return_logits: Logical, default TRUE. If TRUE, return raw logits instead of probabilities. Used only in inference mode.

  • softmax_temperature: Float, default 0.9. Temperature for the softmax function. Used only in inference mode.

  • inference_config: An inference_config object. Used only in inference mode.

Returns a tensor. For training mode: predictions of shape (B, test_size, out_dim). For inference mode: logits or probabilities of shape (B, test_size, num_classes) for classification, or predictions of shape (B, test_size, num_quantiles) for regression.

predict_stats(X, y_train, output_type, alphas, embed_with_test, inference_config)

Compute summary statistics from predicted quantiles. Only applicable for regression tasks (max_classes = 0).

  • output_type: Character string or character vector. Supported values: "mean", "variance", "median", "quantiles", "raw_quantiles". If a vector, returns a named list.

  • alphas: Optional numeric vector. Probability levels for quantile output. Default: c(0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9).

Returns a tensor (single output type) or a named list of tensors (multiple output types). Output shapes: "mean", "variance", "median" return (B, test_size); "quantiles" returns (B, test_size, len(alphas)); "raw_quantiles" returns (B, test_size, num_quantiles).

forward_with_cache(X_train, y_train, X_test, return_logits, softmax_temperature, use_cache, store_cache, cache, cache_mode, inference_config)

Forward pass with caching support for efficient inference. Two caching modes are supported:

  • "kv": Cache KV projections from both column embedding and ICL transformer layers. Fastest inference but uses more memory.

  • "repr": Cache column embedding KV projections and row interaction outputs. Uses approximately 24x less memory for the ICL part, at the cost of re-running the ICL transformer.

Exactly one of use_cache or store_cache must be TRUE. When store_cache = TRUE, requires X_train and y_train. When use_cache = TRUE, requires X_test and a populated cache.

  • X_train: Optional tensor of shape (B, train_size, H). Required when store_cache = TRUE.

  • y_train: Optional tensor of shape (B, train_size). Required when store_cache = TRUE.

  • X_test: Optional tensor of shape (B, test_size, H). Required when use_cache = TRUE.

  • return_logits: Logical, default TRUE.

  • softmax_temperature: Float, default 0.9.

  • use_cache: Logical, default FALSE.

  • store_cache: Logical, default TRUE.

  • cache: Optional TabICLCache. If provided, equivalent to setting use_cache = TRUE and store_cache = FALSE.

  • cache_mode: Character string, default "kv". Caching strategy.

  • inference_config: Optional inference_config.

Returns predictions of shape (B, test_size, out_dim), or NULL if store_cache = TRUE and X_test is not provided.

predict_stats_with_cache(X_train, y_train, X_test, output_type, alphas, use_cache, store_cache, cache, cache_mode, inference_config)

Compute summary statistics from predicted quantiles with KV caching. Only applicable for regression tasks (max_classes = 0). Parameters and return value are the same as predict_stats(), with additional caching parameters matching forward_with_cache(). Returns NULL if store_cache = TRUE and X_test is not provided.

has_cache()

Check if a valid cache is stored. Returns a logical value.

clear_cache()

Clear the stored cache. Called for its side effect; returns NULL invisibly.

See also

ColEmbedding, RowInteraction, ICLearning, quantile_to_distribution, TabICLCache, inference_config