TabICL

TabICL is a transformer-based foundation model that uses in-context learning. Instead of traditional training, it learns the mapping between features and targets from a provided training set during a single forward pass.

The Process:

1. fit(X, y): Preprocesses training data, creates multiple transformed dataset views (e.g., by shuffling features), and optionally pre-computes KV caches for the training data (controlled by the kv_cache parameter during initialization) to speed up inference.
2. predict(X): Processes test data and forwards each dataset view through the model. The final prediction is the average across all ensemble members.

Architecture: It uses a three-stage Transformer architecture:

• Column-wise Transformer: Embeds each feature.
• Row-wise Transformer: Aggregates features into row representations.
• Dataset-wise Transformer: Performs in-context learning over training and test samples.

TabICL supports zero-shot time series forecasting via the TabICLForecaster. This requires the forecast extra.

Installation:

pip install tabicl[forecast]

Key Parameters:

• max_context_length: Maximum historical timesteps to use as context.
• temporal_features: Features to extract (e.g., ["index", "datetime", "periodic"] or custom TimeTransform instances).
• point_estimate: Method for point prediction ("mean" or "median").
• tabicl_config: Configuration object passed to the underlying TabICLRegressor.

Usage Example:

import pandas as pd
from tabicl import TabICLForecaster
from tabicl.forecast import TimeSeriesDataFrame, plot_forecast

# Load data
df = pd.read_csv("data.csv", parse_dates=["timestamp"])
data = TimeSeriesDataFrame.from_data_frame(df)

# Split and prepare
prediction_length = 96
train_data, test_data = data.train_test_split(prediction_length)

# Forecast
forecaster = TabICLForecaster(max_context_length=10240)
pred_df = forecaster.predict_df(context_df, prediction_length=prediction_length)

# Visualize
fig, axes = plot_forecast(context_df=context_df, pred_df=pred_df, test_df=test_df)

import pandas as pd
from tabicl import TabICLForecaster
from tabicl.forecast import TimeSeriesDataFrame, plot_forecast

df = pd.read_csv(
    "https://autogluon.s3.amazonaws.com/datasets/timeseries/australian_electricity_subset/test.csv",
    parse_dates=["timestamp"],
)
data = TimeSeriesDataFrame.from_data_frame(df)

prediction_length = 96
selected_items = data.item_ids[:2]
train_data, test_data = data.train_test_split(prediction_length)

context_df = train_data.reset_index()
context_df = context_df[context_df["item_id"].isin(selected_items)]
test_df = test_data.reset_index()
test_df = test_df.groupby("item_id").tail(prediction_length)

forecaster = TabICLForecaster(max_context_length=10240)
pred_df = forecaster.predict_df(context_df, prediction_length=prediction_length)
fig, axes = plot_forecast(context_df=context_df, pred_df=pred_df, test_df=test_df)

Pre-training is available for both TabICLv1 and TabICLv2. The easiest method is to use the stage scripts located in the scripts folder, which wrap python -m tabicl.train with torchrun.

Pre-training Workflow

1. Open the desired stage script (e.g., scripts/train_v2_clf_stage1.sh).
2. Adjust NUM_GPUS, --n_jobs, and placeholder checkpoint paths for your hardware.
3. Run the three stages in order (Stage 1 $\rightarrow$ Stage 2 $\rightarrow$ Stage 3). For TabICLv2 classifier, use train_v2_clf_stage1.sh, stage2.sh, and stage3.sh.

Key Training Options

• Data Generation: By default, tabicl.train generates synthetic prior datasets on the fly. Alternatively, you can pre-generate datasets to disk using python -m tabicl.prior --save_dir /path/to/dir --num_batches 100000 ... and load them via the --prior_dir flag.
• Loss Functions: Supports classification (cross-entropy) and quantile regression (use --regression_method quantile for pinball loss).
• Optimizers: Supports AdamW (default) and Muon (use --muon True).
• Weight Decay: Cautious weight decay is available via --use_cautious_wd, though it is disabled by default in the provided v2 scripts to match the original released checkpoints.

# Example for TabICLv2 classifier pre-training stages
bash scripts/train_v2_clf_stage1.sh
bash scripts/train_v2_clf_stage2.sh   # loads the stage-1 checkpoint
bash scripts/train_v2_clf_stage3.sh   # loads the stage-2 checkpoint

While TabICL works zero-shot by default, you can use FinetunedTabICLClassifier or FinetunedTabICLRegressor to specialize the model on a specific dataset using a full PyTorch training loop (AdamW, cosine-with-warmup, early stopping).

First, install the required dependencies:

pip install tabicl[finetune]

Usage Example:

from tabicl import FinetunedTabICLClassifier

clf = FinetunedTabICLClassifier(
    epochs=50,
    learning_rate=1e-5,
    n_estimators_finetune=2,
    n_estimators_validation=2,
    n_estimators_inference=8,
    early_stopping=True,
    patience=10,
    eval_metric="roc_auc",
    random_state=0,
    verbose=True,
)

clf.fit(X_train, y_train, X_val=X_val, y_val=y_val, output_dir="./ckpts")
y_pred = clf.predict(X_test)

Loading a fine-tuned model: The checkpoints saved in output_dir are compatible with the standard zero-shot estimators:

from tabicl import TabICLClassifier
clf = TabICLClassifier(model_path="ckpts/best.ckpt")
clf.fit(X_train, y_train)

Multi-GPU training: Use torchrun to enable multi-GPU fine-tuning:

torchrun --nproc-per-node=2 finetune_script.py

Built-in Preprocessing

TabICL accepts pandas DataFrames or numpy arrays and automatically applies:

• Categorical Encoding: Detects and ordinal encodes string, object, category, and boolean types. For numpy arrays, it uses the array's datatype. Integers are treated as numerical.
• Missing Values: Creates a separate category for missing categorical features and performs mean imputation for missing numerical values (encoded as NaN).
• Cleaning: Outlier detection/removal, feature scaling, normalization, and feature shuffling for ensemble diversity.

Advanced Preprocessing with skrub

For complex, heterogeneous, or "dirty" real-world data, it is recommended to use the skrub library. You can integrate skrub.TableVectorizer into a scikit-learn pipeline with TabICLClassifier to handle diverse data types (text, datetime, etc.) before passing them to the model.

# Installation: pip install skrub -U
from skrub import TableVectorizer
from tabicl import TabICLClassifier
from sklearn.pipeline import make_pipeline

pipeline = make_pipeline(
    TableVectorizer(low_cardinality="passthrough"),  # Automatically handles various data types
    TabICLClassifier()
)

pipeline.fit(X_train, y_train)  # X should be a DataFrame
predictions = pipeline.predict(X_test)