TransTab: Transferable Tabular Learning Framework

repository·main·Indexed 42 days ago

https://github.com/ryanwangzf/transtab

A flexible framework for transfer learning and zero-shot learning on tabular data using Transformers. TransTab handles heterogeneous datasets with varying column sets by incorporating column names into feature processing. It supports categorical, numerical, and boolean feature types and utilizes Vertical Partition Contrastive Learning (VPCL) for pretraining across multiple datasets. Key functionality includes building classifiers, regressors, and encoders via `transtab.build_classifier`, `transtab.build_regressor`, and `transtab.build_encoder`.

Tokens
15.9K
Snippets
59
Records
75
Agent score
50%

What's inside TransTab

  1. Supported tabular prediction scenarios

    main

    Unlike traditional methods that require fixed-column tables, transtab supports variable-column tables across several machine learning scenarios:

    • Supervised learning: Standard training and prediction on tables with identical columns.
    • Transfer learning: Leveraging multiple labeled tables that partially share columns to enhance models for each table.
    • Incremental learning: Updating an existing model to handle a table that has grown with new columns.
    • Table Pretraining: Pretraining models on many tables with distinct columns and identifiers for a target task.
    • Zero-shot inference: Building a model for an unseen table that only has partial overlaps with training tables.
  2. Understand TransTab's feature processing

    main

    TransTab achieves flexibility by treating column names as semantic indicators. This allows the model to handle tables with different structures (different columns) by embedding the relationship between a column's name and its value.

    Feature Types Supported:

    • Categorical: The model concatenates the column name and the feature value (e.g., column_name is value).
    • Numerical: The model tokenizes and embeds the column name, then scales that embedding by the numerical feature value.
    • Boolean (bin_cols): The model tokenizes and embeds the column name. If the feature value is 0, the embedding is dropped; otherwise, it is passed to the encoder.
  3. Constraint: Label class consistency in supervised learning

    main
    When using TransTab in a pure supervised learning setting, all datasets provided must have the same number of label classes. This is because the classifier maintains a single classification head for both training and prediction. For example, if training on multiple datasets, all must be binary classification or all must be multiclass with the same number of classes.
  4. How Vertical Partition Contrastive Learning (VPCL) works

    main

    TransTab uses Vertical Partition Contrastive Learning (VPCL) for pretraining across multiple heterogeneous datasets. This approach avoids the pitfalls of naive supervised learning (where different datasets might have conflicting label definitions) by using a contrastive objective.

    The VPCL Process:

    1. Vertical Partitioning: Each raw sample is split vertically into several partitions. Each partition acts as a separate sample for contrastive learning.
    2. Contrastive Pairs: Partitions derived from raw samples sharing the same label are treated as positive pairs. Partitions from samples with different labels are treated as negative pairs.

    This method expands the number of available training pairs and is more robust than adding task-specific classifiers for every dataset.

  5. Quickstart: Basic Classification and Regression

    main

    TransTab provides a high-level API to load data, build models, train them, and make predictions in a few lines of code.

    To use the framework:

    1. Load data using transtab.load_data(dataset_name).
    2. Build a model using transtab.build_classifier or transtab.build_regressor.
    3. Train the model using transtab.train.
    4. Predict using transtab.predict.
    import transtab
    
    # load dataset by specifying dataset name
    allset, trainset, valset, testset, cat_cols, num_cols, bin_cols \
         = transtab.load_data('credit-g')
    
    # build classifier
    model = transtab.build_classifier(cat_cols, num_cols, bin_cols)
    
    # build regressor
    # model = transtab.build_regressor(cat_cols, num_cols, bin_cols)
    
    # start training
    transtab.train(model, trainset, valset, **training_arguments)
    
    # make predictions, df_x is a pd.DataFrame with shape (n, d)
    # return the predictions ypred with shape (n, 1) if binary classification;
    # (n, n_class) if multiclass classification.
    ypred = transtab.predict(model, df_x)
  6. Quickstart: Basic usage of transtab

    main

    You can perform tabular prediction tasks in a few lines of code using transtab. The workflow typically involves loading a dataset, building a model (classifier or regressor), training the model, and making predictions. transtab is designed to handle variable-column tables, supporting supervised learning, transfer learning, incremental learning, table pretraining, and zero-shot inference.

    import transtab
    
    # 1. Load dataset by specifying dataset name
    allset, trainset, valset, testset, cat_cols, num_cols, bin_cols \
        = transtab.load_data('credit-g')
    
    # 2. Build classifier
    model = transtab.build_classifier(cat_cols, num_cols, bin_cols)
    
    # 3. Start training
    transtab.train(model, trainset, valset, **training_arguments)
    
    # 4. Make predictions
    # df_x is a pd.DataFrame with shape (n, d)
    # Returns ypred with shape (n, 1) for binary or (n, n_class) for multiclass
    ypred = transtab.predict(model, df_x)
  7. Fine-tune a pretrained TransTab model

    main

    After completing pretraining, you can load the pretrained weights into a new classifier for fine-tuning on a target dataset.

    1. Use transtab.build_classifier(checkpoint='PATH_TO_CHECKPOINT') to load the model.
    2. Use model.update({'cat': cat_cols, 'num': num_cols, 'bin': bin_cols}) to map the model to the specific column structure of your target dataset.
    # load the pretrained model and finetune on a target dataset
    allset, trainset, valset, testset, cat_cols, num_cols, bin_cols \
        = transtab.load_data('credit-approval')
    
    # build transtab classifier model, and load from the pretrained dir
    model = transtab.build_classifier(checkpoint='./checkpoint')
    
    # update model's categorical/numerical/binary column dict
    model.update({'cat':cat_cols,'num':num_cols,'bin':bin_cols})
  8. Install TransTab from GitHub

    main

    To install TransTab, first ensure you have the correct pytorch version installed following the official PyTorch guide.

    Note: The PyPI version is currently not maintained. Install directly from the GitHub repository using pip.

    pip install git+https://github.com/RyanWangZf/transtab.git
  9. Conduct Contrastive Pretraining on multiple tables

    main

    You can perform contrastive pretraining using multiple distinct datasets simultaneously.

    1. Pass a list of dataset names to transtab.load_data(dataname_list).
    2. Create a contrastive learner using transtab.build_contrastive_learner(cat_cols, num_cols, bin_cols, supervised=True).
    3. Train the model using transtab.train, ensuring you pass the collate_fn returned by the learner.
    # load from multiple tabular datasets
    dataname_list = ['credit-g', 'credit-approval']
    allset, trainset, valset, testset, cat_cols, num_cols, bin_cols \
         = transtab.load_data(dataname_list)
    
    # build contrastive learner, set supervised=True for supervised VPCL
    model, collate_fn = transtab.build_contrastive_learner(
        cat_cols, num_cols, bin_cols, supervised=True)
    
    # start contrastive pretraining training
    transtab.train(model, trainset, valset, collate_fn=collate_fn, **training_arguments)
  10. Use TransTab for tabular prediction

    main

    TransTab is a flexible framework for transfer learning and zero-shot learning on tabular data. Unlike traditional models that require fixed-structure tables, TransTab encodes arbitrary input tables by incorporating column names into the feature processing, allowing it to handle heterogeneous datasets and varying column sets.

    To use TransTab, you can load datasets by name, build a classifier or regressor, and then train or predict. The model handles three feature types:

    • Categorical: Concatenates column names and features (e.g., gender is male).
    • Numerical: Tokenizes and embeds column names, then multiplies the embeddings by the feature value.
    • Boolean: Tokenizes and embeds column names to decide whether to pass the embedding based on the feature value.
    # Example workflow sketch
    # 1. Load datasets
    datasets = ['dataset_name_1', 'dataset_name_2']
    
    # 2. Build a model (Classifier or Regressor)
    # model = transtab.build_classifier(cat_cols, num_cols, bin_cols)
    model = transtab.build_regressor(cat_cols, num_cols, bin_cols)
    
    # 3. Start training
    # model.train(datasets)
    
    # 4. Make predictions
    # df_x is a pd.DataFrame with shape (n, d)
    # ypred shape: (n, 1) for binary; (n, n_class) for multiclass
    ypred = model.predict(df_x)