TabNet

Overview

TabNet is a deep learning model for tabular data that chooses which features to look at, one decision step after another. At each step it builds a sparse mask over the input columns, reads only the selected ones, and adds what it learns to a running summary. The result is a network that predicts well and also shows which features it used for each row.

Its main selling points are built-in feature selection and built-in interpretability. Unlike a tree ensemble, it is trained end to end with gradient descent, so it can be pre-trained with a self-supervised task and plugged into larger neural systems. TabNet came out before the current wave of tabular transformers, and many later models, including those described at /model/architecture, either borrow its ideas or measure themselves against it.

What it is

TabNet stands for Tabular Network. It is a neural architecture for classification and regression on tables. It takes raw features directly, with no heavy manual feature engineering, and uses a chain of decision steps. Each step has two small modules.

• An attentive transformer that decides which features to use at this step. It produces a sparse mask over the columns.
• A feature transformer that processes the selected features and produces both a piece of the final decision and a signal to guide the next step.

Think of it as a detective who reads a case file in several passes. On each pass the detective picks a few documents to study, notes down what they reveal, and uses those notes to decide which documents to read next. After a few passes the notes are combined into a verdict.

What it aims to give you

• Sparse selection. Only a few features matter at each step, which cuts noise.
• Instance-wise interpretability. The masks show which features mattered for each individual row.
• Global interpretability. Summing the masks across steps and rows gives overall feature importance.
• End-to-end learning. No separate feature selection stage is needed.
• Self-supervised pre-training. The model can learn from unlabeled rows by reconstructing masked features.

Typical use

TabNet is used when you want a neural model for a tabular task and also want a quick view of feature use. It is available in open-source form for PyTorch, which makes it easy to try beside scikit-learn style workflows.

Where it came from

TabNet was introduced by Arik and Pfister in 2019 in the paper "TabNet: Attentive Interpretable Tabular Learning". It was later presented at AAAI in 2021. The authors worked on applying deep learning to structured data in practical settings, where trees were the default.

At that time deep learning had transformed images, speech and text, but tabular data was still the home of gradient boosted trees such as XGBoost, LightGBM and CatBoost. The paper argued that trees have two properties worth copying: they select relevant features and they split the space in a few simple decisions. TabNet tries to get these properties inside a neural network by using sequential attention and sparse masks.

Ideas it built on

• Attention. The concept of learning to focus on parts of the input, as used in Transformers.
• Sparsemax. A softmax alternative that can output exact zeros, proposed in earlier work on sparse attention.
• Gated linear units. Gating layers used in language modeling.
• Batch normalization with ghost batches. A trick to stabilize training with large batches.
• Decision trees. The inspiration for the step-by-step selection.

What came after

TabNet had a real influence on the field. It was widely adopted through community libraries, and a large number of tutorials use it. Later work showed that on many benchmark datasets, tuned trees or simpler neural models such as FT-Transformer matched or exceeded it, which led to a more careful view of its strengths. It remains valuable for its interpretability story and its masked-feature pre-training. Related works like SAINT and XTab continued the trend of pre-training for tables.

How it works

The best way to understand TabNet is to follow one row through the network.

Step 0: Normalize the input

The raw feature vector is passed through batch normalization. Categorical features are first turned into embeddings or integer-coded and embedded. Numeric features are used as they are, after scaling.

Step 1: Decide what to look at

The first decision step has no previous information, so the attentive transformer uses a neutral starting point. It outputs a mask M with one value per feature. Values are non-negative and sum to one across features, and many are exactly zero because of the sparsemax function.

Everyday analogy: you walk into a warehouse of 100 boxes and your instructions say "open at most five boxes now". The mask is your list of five boxes with a weight for each.

Step 2: Read the selected features

The mask multiplies the input features, so unselected features become zero. The masked vector goes into the feature transformer.

Step 3: Process with the feature transformer

The feature transformer is a small stack of fully connected layers with batch normalization and gated linear units. Some layers are shared across all decision steps, and some are specific to each step. This sharing keeps the model compact and lets it learn common patterns once.

Its output is split in two parts.

• One part goes to the decision output for this step, after a ReLU.
• The other part goes to the attentive transformer of the next step, to help it decide what to read next.

Step 4: Decide what to look at next

The next attentive transformer receives the signal from the last step and a prior scale term. The prior scale tracks how much each feature has already been used. If a feature was used heavily in earlier steps, its prior shrinks, which pushes the model to look at other features. A parameter called the relaxation factor controls how strongly this applies. A value of one means each feature is used about once overall. Larger values let features be reused.

Analogy: after opening five boxes, you cross them off. The next pass favors boxes you have not opened, though you may reopen one if it was really useful and you have allowed reuse.

Step 5: Combine the steps

The decision outputs from all steps are added together. A final linear layer maps the sum to the prediction. Typical models use between three and ten steps.

Step 6: Aggregate masks for explanation

For each row, TabNet weights each step's mask by how much that step contributed to the decision, then sums over steps. The result is a feature importance vector for that row. Averaging over many rows gives global importance. Because masks are sparse, the explanation lists a handful of features rather than all of them.

Step 7: Train with a sparsity regularizer

The loss has two parts. One is the task loss, such as cross-entropy. The other is an entropy-based penalty on the masks that encourages them to be sparse. A coefficient sets the balance. A larger value gives sparser masks, but too large can hurt accuracy.

Step 8: Optional self-supervised pre-training

TabNet has an encoder-decoder form for pre-training. A random subset of feature values is hidden. The encoder reads the visible ones, and the decoder tries to reconstruct the hidden values. This works like filling in blanks in a table. After pre-training, the encoder is reused for supervised training. The paper showed that this helps when labeled data is limited and unlabeled data is abundant.

Why ghost batch normalization

Ghost batch normalization splits a large batch into small virtual batches and normalizes each one separately. This keeps the noise level of small batches while still using large batches for speed. It is a useful trick because TabNet is sensitive to normalization statistics.

Key concepts and terms

A worked example

This example uses a common open-source implementation for PyTorch. It shows training on a classification task and reading feature importance.

import numpy as np
from pytorch_tabnet.tab_model import TabNetClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score

# X: numpy array of features, y: binary labels
X_train, X_valid, y_train, y_valid = train_test_split(
    X, y, test_size=0.2, random_state=0, stratify=y)

clf = TabNetClassifier(
    n_d=16, n_a=16,          # width of decision and attention parts
    n_steps=4,               # number of decision steps
    gamma=1.5,               # relaxation factor for feature reuse
    lambda_sparse=1e-4,      # sparsity regularization weight
    seed=0,
)

clf.fit(
    X_train, y_train,
    eval_set=[(X_valid, y_valid)],
    eval_metric=["auc"],
    max_epochs=200,
    patience=20,             # early stopping
    batch_size=1024,
    virtual_batch_size=128,  # ghost batch size
)

pred = clf.predict_proba(X_valid)[:, 1]
print("AUC:", roc_auc_score(y_valid, pred))

# Global feature importance
print(clf.feature_importances_)

# Per-row explanation: masks for each step
explain_matrix, masks = clf.explain(X_valid)

Reading the explanation

explain_matrix has one row per sample and one column per feature. A row shows how much each feature contributed to that prediction, summed across steps. masks is a dictionary that gives the mask for each step separately, so you can see that step one focused on age and balance while step three looked at tenure.

Self-supervised pre-training in the same library

from pytorch_tabnet.pretraining import TabNetPretrainer

pretrainer = TabNetPretrainer(n_d=16, n_a=16, n_steps=4, mask_type="entmax")
pretrainer.fit(
    X_train=X_unlabeled,
    eval_set=[X_valid],
    pretraining_ratio=0.8,   # fraction of features to hide
    max_epochs=100,
)

clf = TabNetClassifier(n_d=16, n_a=16, n_steps=4)
clf.fit(X_train, y_train,
        eval_set=[(X_valid, y_valid)],
        from_unsupervised=pretrainer)

A small numeric picture

Say a row has six features: age, income, balance, tenure, region and plan. With four steps and a sparse mask, step one might put weight 0.7 on balance and 0.3 on income. Step two, helped by the prior scale, might pick tenure and plan. Step three might return to income, since a relaxation factor above one allows reuse. Step four might focus on age. The aggregated explanation could then read: balance 0.35, income 0.30, tenure 0.15, plan 0.10, age 0.10, region 0. That is a short, readable account of the decision.

Tuning example

A reasonable start is n_d = n_a between 8 and 64, n_steps between 3 and 7, gamma between 1.0 and 2.0, and lambda_sparse between 0 and 0.01. Search them with Optuna while keeping the learning rate near 0.02 with a decay schedule.

Sparsemax compared with softmax

Both functions turn a list of scores into weights that add up to one. The difference shows in how they treat low scores.

Softmax never reaches zero, so every feature always leaks a little into the result. Sparsemax can drop weak features entirely. That is why TabNet masks can say "this feature was not used at all in this step", which makes explanations short and honest.

A code sketch of one decision step

import torch
import torch.nn as nn

class TinyStep(nn.Module):
    def __init__(self, n_features, n_d, n_a):
        super().__init__()
        self.att = nn.Linear(n_a, n_features)
        self.bn = nn.BatchNorm1d(n_features)
        self.feat = nn.Sequential(
            nn.Linear(n_features, 2 * (n_d + n_a)),
            nn.BatchNorm1d(2 * (n_d + n_a)),
        )
        self.n_d, self.n_a = n_d, n_a

    def forward(self, x, a_prev, prior):
        scores = self.bn(self.att(a_prev)) * prior
        mask = torch.softmax(scores, dim=-1)       # sparsemax in the real model
        h = self.feat(mask * x)
        h = h[:, : self.n_d + self.n_a] * torch.sigmoid(h[:, self.n_d + self.n_a:])
        d, a = h[:, : self.n_d], h[:, self.n_d:]
        prior = prior * (1.5 - mask)                # gamma = 1.5
        return torch.relu(d), a, mask, prior

This sketch uses softmax to stay short. The real model uses sparsemax or entmax and a deeper feature transformer with shared and step-specific layers. It does show the main flow: mask from the previous step's signal and the prior, masked input, gated processing, a decision part and an attention part, and an updated prior.

How the pieces interact

1. The prior starts at all ones, so every feature is available.
2. The mask for step one comes from the attention signal and the prior.
3. After step one, the prior is multiplied by gamma - mask. Heavily used features get a smaller prior.
4. Step two's mask is shaped by the new prior, so it is nudged toward unused features.
5. The loop continues for the chosen number of steps.
6. The decision parts are added up, and the final layer produces the answer.

A table of how the settings change behavior

What interpretability can and cannot show

Mask weights tell you which input columns the network read at each step. They do not tell you whether changing a column would change the outcome in the real world. Two correlated columns may share the mask or swap between runs. A column that is useful only in combination with another may show a small weight. So the explanation is a useful guide for debugging and review, but any claim about cause needs a separate test, for example a controlled experiment or a causal analysis.

Comparing TabNet masks with tree importance

Typical pre-training setup

• Use all available rows, labeled or not.
• Hide a large share of features, often between half and four fifths of them, at random in each row.
• Train the encoder-decoder to reconstruct the hidden values, with the loss normalized per feature so that wide-ranging columns do not dominate.
• Keep the encoder and discard the decoder, then fine-tune with labels.

Choosing the number of steps

Start with three or four. Raise the count only if validation scores rise and explanations stay readable. Many tables need only a few steps, because each step already reads several features, and extra steps mostly add cost.

A short evaluation plan

1. Fit a tuned tree ensemble and record its validation and test scores.
2. Fit TabNet with default settings, then with a small search.
3. Compare scores, training time and inference time side by side.
4. Inspect the global importance from both models and note where they agree.
5. Check the stability of TabNet importance across three or more seeds.
6. Decide whether the explanation benefit is worth any accuracy gap.

Regression and multi-class notes

• For regression, scale the target and invert the scaling at prediction time.
• For multi-class tasks, watch per-class recall, since rare classes can be ignored by a sparse model.
• For imbalanced binary tasks, use class weights or a sampler and evaluate with precision-recall metrics.

Strengths and limitations

Strengths

• Built-in feature selection. The model learns to ignore irrelevant columns.
• Interpretability by design. Masks give per-row and global views of feature use.
• End-to-end gradient training. It fits into neural pipelines and can share parts with other networks.
• Self-supervised pre-training. Masked reconstruction helps when labels are scarce.
• Handles raw features. Little feature engineering is required.
• Compact. It is often smaller than a large transformer and is quick to train on mid-size tables.
• Good on some data types. It can do well on data with many weak features and some structure.

Limitations

• Not always the most accurate. Later comparisons often found tuned trees or other neural models ahead of it on many datasets.
• Sensitive to settings. Results depend on the relaxation factor, sparsity weight, step count and batch settings.
• Training can be unstable. Without careful normalization and learning-rate scheduling it can diverge or train slowly.
• Mask explanations are model-internal. The masks show what the network read, which is not the same as causal importance.
• Fixed schema. The model is tied to the columns it was trained on and cannot read a new table without retraining.
• Missing values need handling. The basic version expects imputed inputs.
• Slower than trees at inference. Especially for CPU-only serving at high volume.

How it compares with alternatives

If you want the least effort and strong accuracy, start with a tree ensemble. If you want a neural model with a feature-use story, TabNet is a fair option. If you want a modern deep baseline with simpler training, many people try FT-Transformer first.

Role in tabular foundation models

TabNet is a per-dataset model, not a foundation model. Still, it contributes three ideas that matter for models that learn from raw tables: instance-wise feature selection, masked-feature self-supervision and a design that works on mixed raw input.

Row and column structure

TabNet treats a row as a flat vector and gives each column a fixed slot. It does not model the relationship between rows. Its column handling, however, is selective: it learns which columns matter for which row. Foundation models need a more flexible version, where the model chooses columns by attention that adapts to any schema, which is the role attention plays in FT-Transformer and in the architectures at /model/architecture.

Mixed types

Categorical columns get embeddings and numeric columns are scaled. This is simple and works. Foundation models extend the idea with type-aware tokenizers that generalize to unseen columns, as described for XTab and TabSTAR.

Missing values

The masking idea is a close cousin of missingness. In pre-training, TabNet hides features at random and learns to reconstruct them. At inference, a missing value can be treated like a hidden feature, and the encoder can still produce a representation. This links TabNet to Masked cell modeling and to Imputation networks. In the basic supervised version, missing values should still be imputed or flagged.

No column names

TabNet does not use column names at all. Slots are positional. Feature importance outputs are indexes that you map back to names yourself. This is a limit for foundation models, which should use names when they exist and still work when they do not. Text-aware approaches such as TabSTAR use names to build better tokens, while TabNet-style models would need to be retrained for any schema change.

Pre-training on synthetic or real tables

TabNet pre-trains only on the same table it will later classify. There is no transfer to other tables. Cross-table pre-training, as in XTab, and pre-training on synthetic tables, as in TabPFN, go much further. What TabNet shows is that masked reconstruction is a good unlabeled signal for tables, and this idea continues in newer work. Compare results across models at /model/benchmarks and see the research context at /model/research.

In-context learning

TabNet learns by gradient descent on the target data and does not use labeled rows as a prompt. Its sequential attention is over features, not over examples. In-context learners such as TabPFN and TabICL move the learning step into the forward pass. TabNet's steps could be seen as a small, fixed-length reasoning chain, but they never look at other rows.

Continual learning

A TabNet model can be fine-tuned when new data arrives. Without safeguards it may forget earlier patterns. Methods such as Elastic weight consolidation and Experience replay apply directly, because they act on weights and training samples. The sparse masks add one more handle: you can watch the average mask over time to see if the model's focus drifts.

Imputation and anomaly detection

The decoder in the pre-training mode reconstructs hidden features, which is imputation. The reconstruction error on a row, compared with typical error, gives an anomaly signal. See Anomaly detection and Isolation forest for other approaches. Mask patterns can also signal anomalies: a row for which the model needs unusual features may be unusual itself.

Sector and semantic classification

For tasks that label a row or a table by domain, the feature mask gives a built-in clue about which columns drove the label. When columns have meaningful names, a human can read the explanation. For a model that works across schemas, Semantic sector heads show how a classifier on top of learned table representations can assign domain labels, and column importance masks can help audit them.

Practical guidance and best practices

When TabNet is a good fit

• You want a neural tabular model with explanations you can show a reviewer.
• You have many columns and expect only some to matter.
• You can use unlabeled rows for pre-training.
• You plan to combine the model with other neural components.

When to prefer something else

• The table is small. Trees or in-context models are safer.
• Pure accuracy is the only goal and compute is tight. Tuned trees often win.
• You need to serve millions of requests on CPU at tiny latency.

Data preparation

1. Split first, then fit scalers and encoders on the training split only.
2. Scale or quantile-transform numeric columns.
3. Label-encode categorical columns and tell the library which columns are categorical, with their cardinalities.
4. Impute missing values or add flags.
5. Remove ID columns and leakage.

Hyperparameters that matter

Training tips

• Use a larger batch than you would for a tree-like experiment, with ghost batch normalization.
• Use a learning-rate schedule with decay.
• Use early stopping with patience tied to the validation metric.
• Run several seeds and average, because results vary.
• Start with pre-training only if you have many more unlabeled rows than labeled ones.

Reading explanations safely

• Combine mask-based importance with permutation importance or SHAP to cross-check.
• Compare feature use for correct and incorrect predictions.
• Look for stable patterns across seeds. If importance changes sharply between seeds, do not trust any single run.
• Never present mask weights as causes.

Deployment notes

• Save the preprocessing, the category maps and the model together.
• Check that feature order at serving matches training.
• Log masks for a sample of requests to monitor drift in feature use.
• Track experiments with MLflow or a similar tool.

Common mistakes

• Using default settings and stopping there. TabNet is sensitive, and defaults often underperform.
• Comparing to untuned trees or to a poorly tuned TabNet. Both sides need fair budgets.
• Setting sparsity too high. The masks collapse onto a few features and accuracy drops.
• Setting sparsity to zero and expecting interpretability. Dense masks explain little.
• Forgetting to declare categorical columns. The library will treat them as numbers.
• Skipping scaling. Batch norm helps, but extreme values still hurt.
• Trusting masks as causal. They show where the network looked.
• Leaking test data into pre-training. Keep a clean test set out of the unlabeled pool.
• Changing column order at serving. Slots are positional, so order must match training.
• Single-seed conclusions. Neural training varies between runs.

Frequently asked questions

What is TabNet used for?

TabNet is used for classification and regression on tabular data when you want a neural network with built-in feature selection and explanations. It reads the inputs in several decision steps, each with a sparse mask, and it can be pre-trained on unlabeled rows.

How does TabNet select features?

At each decision step an attentive transformer produces a mask over the columns using a sparsemax function, which can give exact zeros. The mask is multiplied with the input so only a few features are read. A prior scale discourages reusing the same features too often.

Is TabNet interpretable?

It is interpretable in a practical sense, because its masks show which features each step used for each row, and they can be summed into global importance. This is a view of the network's behavior and not a proof of cause, so it should be cross-checked with other methods.

Is TabNet better than XGBoost?

Usually not on accuracy. Later comparisons found tuned gradient boosted trees to be strong, and TabNet is more sensitive to settings. It can still be a good choice when you need a neural model, self-supervised pre-training or mask-based explanations.

What is sparsemax?

Sparsemax is a function similar to softmax that turns scores into weights summing to one, but it can set some weights exactly to zero. TabNet uses it so that each decision step looks at a small set of features rather than spreading attention across all of them.

Can TabNet handle categorical features and missing values?

Categorical features are handled with embeddings, which you must declare in the setup. Missing values are not handled well by the basic supervised model, so impute them or add indicator features. The pre-training mode, which reconstructs hidden features, is conceptually close to imputation.

Does TabNet need a lot of data?

It tends to work better with thousands of rows or more. On very small tables it can overfit or train unstably, and simple models or in-context learners are safer. Pre-training on unlabeled rows can help when labels are limited but raw data is plentiful.

How is TabNet different from a transformer?

TabNet uses sequential attention over features through masks, with fully connected processing layers, and it does not look at relationships between every pair of features. A transformer like FT-Transformer uses self-attention across all feature tokens at once, which is more flexible but costs more on wide tables.

Is TabNet a tabular foundation model?

No. It is trained on one dataset and tied to its columns. It contributes ideas that foundation models reuse, such as masked-feature self-supervision and learned feature selection, but it cannot read a new table without retraining.

What do the steps in TabNet mean?

Each step is one round of choosing features, processing them and adding a contribution to the final decision. More steps allow more complex combinations and more feature reuse, but they raise cost and the risk of overfitting. Most models use between three and ten.

Glossary

• Attentive transformer: the module that builds the feature mask at each step.
• Batch normalization: a layer that rescales activations over a batch for stable training.
• Decision step: one round of feature selection and processing in TabNet.
• Encoder-decoder: a pair of networks, one compressing the input and one reconstructing it.
• Entmax: a family of functions between softmax and sparsemax.
• Feature importance: a score of how much each feature matters.
• Feature transformer: the module that processes the masked inputs.
• Gated linear unit: a layer where one half of the output scales the other half.
• Ghost batch normalization: normalization within small virtual sub-batches.
• Instance-wise: computed separately for each row.
• Mask: a vector of weights used to select features.
• Prior scale: a running record that reduces the weight of already used features.
• Relaxation factor: a parameter that lets features be reused across steps.
• Self-supervised learning: learning from signals created from the data itself. See Self-supervised learning.
• Sparsemax: a softmax-like function that can output zeros.

Summary

TabNet is a neural tabular model that reads features through a chain of sparse attention masks. Each decision step picks a few columns, processes them, and adds to a running decision. The masks double as explanations, both for single rows and across the dataset, and an optional masked-reconstruction stage lets the model learn from unlabeled rows.

It is best seen as a useful, interpretable neural option rather than a guaranteed accuracy leader, since tuned trees such as XGBoost and neural baselines like FT-Transformer often match or beat it. Its lasting value to the field is conceptual. Learned feature selection and masked-feature pre-training both carry into modern tabular foundation models such as XTab, TabPFN and TabSTAR. Use it with careful tuning, cross-check its explanations, and compare it fairly with strong baselines. More background is at /model/research, /model/architecture and /model/benchmarks.

Related

It belongs to Related models and baselines in my stack, under Foundation models and applied ML. Others in the same group: