DOCSEARCHbyAlgolia
Submit Docs

TabPFN Extensions

repository·main·Indexed 43 days ago

https://github.com/priorlabs/tabpfn-extensions

Experimental tools to extend the TabPFN foundation model for tabular data (v0.6.2). Includes modules for interpretability, unsupervised learning, Bayesian optimization, and image integration via TabPFNWithImages and ImageTransformer. Provides TabPFNEmbedding for extracting dense sample embeddings and advanced predictive distribution analysis for TabPFNRegressor, including quantiles and credible intervals.

Tokens
14.8K
Snippets
40
Records
55
Agent score
55%

What's inside TabPFN Extensions

  1. Overview of TabEBM for Tabular Data Augmentation

    main

    TabEBM (Tabular Data Augmentation Method with Distinct Class-Specific Energy-Based Models) is a high-performance method designed to improve the performance of predictive models on small tabular datasets.

    Unlike traditional generative methods that learn a single joint or class-conditional distribution, TabEBM creates distinct Energy-Based Models (EBMs) for each class. This allows the system to model class-specific data distributions individually, creating robust energy landscapes even in ambiguous distributions.

    Key characteristics include:

    • No additional training required for downstream models: The synthetic data generated can be applied to any existing predictive model.
    • High performance: Optimized with GPU acceleration, advanced caching, and memory-efficient Stochastic Gradient Langevin Dynamics (SGLD) sampling.
    • TabPFN-v2 Integration: Uses TabPFN-v2 to enhance gradient-based sampling via its energy landscapes.

  2. Overview of available TabPFN extensions

    main

    The tabpfn-extensions repository provides several specialized modules to extend the core TabPFN functionality:

    • interpretability: Explain predictions using SHAP values and feature selection.
    • many_class: Handle classification problems with more classes than the TabPFN checkpoint supports.
    • unsupervised: Includes data generation, imputation, and outlier detection.
    • embedding: Access TabPFN's internal dense sample embeddings.
    • image: Process image columns (via file paths, base64, bytes, or PIL images) by converting them into frozen DINOv3 features.
    • tabebm: Data augmentation using TabPFN-based Energy-Based Models.
    • pval_crt: Statistical feature relevance testing using p-values.
    • bayesian_optimization: Uses TabPFN as a surrogate model with differentiable Expected Improvement.

  3. How to choose between Mean, Median, and Mode

    main

    The choice of point estimate depends on your loss function and the shape of the predictive distribution:

    OutputBest forNotes
    MeanMSE-style metricsOptimal under squared loss; can fall in empty troughs of multimodal distributions.
    MedianRobustness / MAEMinimizes Mean Absolute Error; robust to heavy tails.
    ModeMultimodal posteriorsThe "most likely" answer. Calculated as the bucket with highest density (probability/width).
    QuantilesRisk-aware decisionsBest for calibrated intervals and understanding uncertainty.

  4. Configure TabPFN backend options

    main

    Most extensions in tabpfn-extensions support two different backends for inference. You can choose the one that fits your infrastructure requirements:

    1. TabPFN Package: A full PyTorch implementation for local inference.
      • Install via: pip install tabpfn
    2. TabPFN Client: A lightweight API client for cloud-based inference.
      • Install via: pip install "tabpfn-client>=0.2.7"

    Important Exceptions: The embedding and bayesian_optimization extensions only work with the local tabpfn package.

    # For local inference
pip install tabpfn

# For cloud-based inference
pip install "tabpfn-client>=0.2.7"

  5. Key Features of TabEBM

    main

    TabEBM provides several optimized capabilities for tabular data workflows:

    • Fast Synthetic Data Generation: Rapid generation of synthetic tabular data with minimal configuration.
    • Memory-Efficient SGLD: Optimized Stochastic Gradient Langevin Dynamics using pre-computed noise tensors to reduce memory overhead.
    • TabPFN-v2 Integration: Seamlessly uses TabPFN-v2 for enhanced gradient-based sampling.
    • Performance Optimizations:
      • Model Caching: An intelligent system that avoids redundant model training across different classes.
      • Vectorized Operations: Optimized tensor operations and reduced device transfers to maximize speed.
      • Stratified Sampling: Smart subsampling for large datasets that maintains class balance.
      • Memory Management: Uses pre-allocated tensors and efficient memory usage patterns.

  6. Configure TabPFNWithImages parameters

    main

    When initializing TabPFNWithImages, you define how to treat specific columns in your input DataFrame X:

    • TabPFNClassifier (or other base estimators): Use categorical_features_indices to specify the positions of categorical columns.
    • image_features_indices: A list of integer positions in the input DataFrame X that contain image data.

    Important: The indices provided to both the base estimator and the TabPFNWithImages wrapper refer to the original frame you pass in. The wrapper internally manages the movement of categorical positions to the expanded frame after image embedding.

  7. How TabPFNEmbedding handles Vanilla vs K-fold embeddings

    main

    Depending on the n_fold parameter, TabPFNEmbedding provides two different embedding strategies:

    1. Vanilla Embeddings (n_fold=0): The model is trained once on the entire training set and used to embed both training and test data. This is simpler but may suffer from distribution shifts because TabPFN attends to training labels.
    2. K-fold Cross-Validation (n_fold >= 2): This is the recommended robust variant. It produces Out-of-Fold (OOF) embeddings for the training data. Each training sample is embedded in a fold where it was held out, matching the condition of unseen data (where no labels are attended to).

    Note on transform(X): The .transform() method always uses the final model trained on the full dataset. It does not return cached OOF embeddings, even if X is the training set. To get OOF embeddings, you must use .fit_transform(X_train, y_train) or access the embedding.train_embeddings_ attribute.

  8. Compute Shapley-based explanations with TabPFN shapiq adapters

    main

    The shapiq integration provides three different adapters for computing Shapley values or interactions, tailored to TabPFN's in-context learning paradigm. Choosing an adapter involves a trade-off between speed and accuracy/methodology.

    1. Remove-and-recontextualize (get_tabpfn_explainer)

    TabPFN is re-fit for every coalition.

    • Pros: High fidelity to the removal paradigm.
    • Cons: Substantially slower because the KV cache cannot be reused across coalitions.

    2. Imputation-based removal (get_tabpfn_imputation_explainer)

    Uses marginal, conditional, or baseline imputation. The training set is fixed, allowing the use of the KV-cache fast path.

    • Requirement: You must construct the model with fit_mode="fit_with_cache" (set before calling .fit()).

    3. Inf-masking removal (get_tabpfn_inf_explainer)

    Sets masked features to +inf, which TabPFN treats as "missing". This uses one forward pass per coalition and utilizes the KV-cache.

    • Requirement: You must construct the model with inference_config={"PASSTHROUGH_INF": True} (requires tabpfn>=8.1.0). This ensures +inf reaches the model instead of being rejected at validation.
    # Example: Imputation-based explainer
clf = TabPFNClassifier(fit_mode="fit_with_cache")
clf.fit(X_train, y_train)
explainer = get_tabpfn_imputation_explainer(model=clf, data=X_train)

# Example: Inf-masking explainer
clf = TabPFNClassifier(
    fit_mode="fit_with_cache",
    inference_config={"PASSTHROUGH_INF": True},
)
clf.fit(X_train, y_train)
explainer = get_tabpfn_inf_explainer(model=clf, data=X_train)

  9. Install tabpfn-extensions

    main

    To install the tabpfn-extensions package with all available extensions, use the following command to install directly from the GitHub repository:

    pip install "tabpfn-extensions[all] @ git+https://github.com/PriorLabs/tabpfn-extensions.git"

    Note: The extensions in this repository are experimental. APIs are subject to change without notice.

  10. Perform feature-level hypothesis testing with TabPFN-CRT

    main

    TabPFN-CRT allows you to test whether a specific feature provides predictive information about a target variable using the Conditional Randomization Test (CRT). Unlike heuristic importance measures like SHAP or permutation importance, this method provides calibrated p-values and formal hypothesis testing by treating TabPFN as a fixed predictive model.

    When to use

    • Validating scientific variables.
    • Auditing black-box model feature importance.
    • Performing feature selection with statistically valid p-values.
    • Working in regulated or high-stakes settings requiring formal hypothesis tests.
    from tabpfn_extensions.pval_crt import tabpfn_crt

# Test a single feature by index
result = tabpfn_crt(
    X,
    y,
    j=3,
    B=200,
    alpha=0.05,
)

print(result["p_value"], result["reject_null"])

  11. Run TabPFN extension examples

    main

    Most examples in the repository are standalone Python scripts that can be executed directly from the command line. Each example is organized into subdirectories corresponding to specific extensions (e.g., interpretability, survival, unsupervised).

    python examples/many_class/many_class_classifier_example.py

  12. Install TabPFN with image support

    main

    To use image-based extensions, install the tabpfn-extensions package with the [image] extra:

    pip install "tabpfn-extensions[image]"

    Note on Model Access: The image encoder uses weights gated on the Hugging Face Hub. You must:

    1. Accept the license on the model page on Hugging Face.
    2. Authenticate locally using hf auth login or by setting the HF_TOKEN environment variable.