API Reference / algorithms / tabular_deep /

saint.py

SAINT: self-attention over features and over rows.

Implements the architecture of Somepalli et al. (2021), SAINT: Improved Neural Networks for Tabular Data via Row Attention and Contrastive Pre-Training. The distinguishing mechanism is intersample attention: after attending across the features of a row, the network attends across the rows of the batch, so a prediction can borrow evidence from neighbouring samples. Without that second stage SAINT would be FT-Transformer.

PyTorch is an optional dependency -- pip install 'tuiml[torch]'. Nothing in this module imports torch until fit is called.

Classes

SAINTClassifier

class algorithms.tabular_deep.saint.SAINTClassifier(_SAINTCore, _DeepTabularClassifierMixin, Classifier)

SAINT: attention across features and across rows of the batch.

SAINT starts where FT-Transformer stops. Attending across the features of a row lets the model see interactions; SAINT adds a second attention stage that runs across the rows of the batch, so the representation of a sample is a function of its neighbours as well as of itself. That is a learned, end-to-end analogue of a nearest-neighbour lookup, and it is what makes SAINT more than a second copy of FT-Transformer.

Overview

  1. Tokenize every feature into a d_token embedding and prepend
[CLS] (the same tokenizer FT-Transformer uses).
  1. Feature attention: a pre-norm Transformer block over the token
axis, relating the columns of one row.
  1. Intersample attention: flatten each row's tokens into a single
vector of width n_tokens * d_token, treat the batch as a sequence of those vectors, and attend over it -- relating whole rows to each other. Reshape back.
  1. Repeat for n_blocks stages and read the [CLS] token.

Theory

Write T^{(i)} \in \mathbb{R}^{t \times d} for the tokens of row i in a batch of size b. Feature attention acts within a row,

T^{(i)} \leftarrow T^{(i)} + \mathrm{MHSA}\big(\mathrm{LN}(T^{(i)})\big),

while intersample attention acts on the flattened matrix Z \in \mathbb{R}^{b \times td} whose rows are \mathrm{vec}(T^{(i)}):

Z \leftarrow Z + \mathrm{MHSA}\big(\mathrm{LN}(Z)\big), \quad \mathrm{Attn}(Q, K, V) = \mathrm{softmax}\! \left(\frac{Q K^{\top}}{\sqrt{d_h}}\right) V .

The two stages differ only in which axis the softmax runs over, and that difference is observable: permuting the other rows of a batch changes a given row's representation under intersample attention, and cannot change it under feature attention alone.

Parameters

d_token
int = 16
Width of each feature token; must be divisible by n_heads. The default is tuned for sub-second fits on toy data; real problems want 32-64 (intersample attention makes SAINT more expensive than FT-Transformer at equal width).
n_blocks
int = 1
Number of (feature attention, row attention) stages. Real use wants 2-6.
n_heads
int = 2
Attention heads per block.
dropout
float = 0.1
Dropout rate in the attention weights and the feed-forward network.
learning_rate
float = 1e-3
AdamW step size.
weight_decay
float = 1e-5
AdamW decoupled weight decay.
batch_size
int = 64
Mini-batch size. It matters more here than in other models: it is the population intersample attention gets to look at, at train and predict time.
n_epochs
int = 60
Passes over the training set; real use wants several hundred.
early_stopping
bool = False
Hold out a validation split and restore the best weights.
validation_fraction
float = 0.15
Fraction held out when early_stopping is enabled.
patience
int = 10
Epochs without validation improvement before training stops.
categorical_features
sequence of int
Column indices holding integer-coded categorical features.
device
{"cpu", "auto", "cuda", "mps"} = "cpu"
Compute device; "auto" trades reproducibility for speed.
random_state
int = 0
Seed for initialisation, shuffling and the validation split.

Attributes

classes_
np.ndarray of shape (n_classes,)
Sorted class labels seen during fit.
network_
torch.nn.Module
The fitted network, rebuilt lazily after unpickling.
feature_mean_, feature_scale_
np.ndarray
Feature standardisation statistics, fitted on train.
loss_curve_
np.ndarray of shape (n_iter_,)
Mean training loss per epoch.
n_iter_
int
Epochs actually run.
n_features_in_
int
Number of features seen during fit.

Notes

Requires PyTorch. Install with pip install 'tuiml[torch]'. The class constructs and introspects without torch; fit raises ImportError naming the install command.

Predictions depend on the batch. Intersample attention is not a per-row function, so a row's prediction depends on which rows share its batch. Predictions are deterministic for a given X and batch_size because inference batches follow input order, but scoring the same row inside a different set of rows can move it. This is inherent to the architecture, not an implementation artifact.

Complexity. Per epoch, O(n m^2 d + n b m d) -- the second term is intersample attention, quadratic in the batch size b. Memory is O(b^2 + b m^2).

When to use. SAINT pays off when rows are informative about each other -- semi-supervised settings, or data with cluster structure the label respects. When rows are genuinely i.i.d. given the features, FTTransformerClassifier gives similar accuracy for less compute.

References

Somepalli2021
Somepalli, G., Goldblum, M., Schwarzschild, A., Bruss, C. B., & Goldstein, T. (2021). SAINT: Improved Neural Networks for Tabular Data via Row Attention and Contrastive Pre-Training. :doi:`10.48550/arXiv.2106.01342`

See Also

FTTransformerClassifier

Feature attention only.
NODEClassifier

Differentiable oblivious trees.
SAINTRegressor

The regression counterpart.

Constructing and inspecting a model needs no torch:

python 
>>> from tuiml.algorithms.tabular_deep import SAINTClassifier
>>> model = SAINTClassifier(d_token=8, n_blocks=2, random_state=0)
>>> model.n_blocks
2
>>> "batch_size" in SAINTClassifier.get_parameter_schema()
True

Fitting requires pip install 'tuiml[torch]'; the example below is a

no-op on an install without it:

python 
>>> import numpy as np
>>> from tuiml.utils.torch_backend import has_torch
>>> rng = np.random.default_rng(0)
>>> X = rng.normal(size=(300, 4))
>>> y = ((X[:, 0] > 0) ^ (X[:, 1] > 0)).astype(int)
>>> if has_torch():
...     model = SAINTClassifier(n_epochs=150, random_state=0).fit(X, y)
...     print(float(model.score(X, y)) > 0.85)
... else:
...     print(True)
True

Methods

get_capabilities (cls) -> List[str]

Return supported capabilities.

get_complexity (cls) -> str

Return complexity analysis.

get_references (cls) -> List[str]

Return academic citations.

SAINTRegressor

class algorithms.tabular_deep.saint.SAINTRegressor(_SAINTCore, _DeepTabularRegressorMixin, Regressor)

SAINT for regression: feature attention plus intersample attention.

The regression counterpart of SAINTClassifier. The stack is identical -- tokenizer, alternating feature and row attention, [CLS] readout -- and only the objective changes: one output unit trained with mean squared error against a standardised target.

Overview

  1. Standardise the target.
  2. Tokenize the features and prepend [CLS].
  3. Alternate feature attention (within a row) and intersample attention
(across the batch) for n_blocks stages.
  1. Read [CLS] through a linear head, then undo the standardisation.

Theory

With Z \in \mathbb{R}^{b \times td} the flattened batch of tokenized rows, intersample attention makes the prediction for row i a weighted combination of every row in the batch:

\hat{y}_i = h\Big(\sum_{j=1}^{b} \alpha_{ij} v(Z_j)\Big), \quad \alpha_{i\cdot} = \mathrm{softmax}\! \left(\frac{q(Z_i) K^{\top}}{\sqrt{d_h}}\right)

which is a learned kernel regression over the batch, trained jointly with the representation it attends over. The loss is mean squared error on the standardised target (y - \mu_y)/\sigma_y.

Parameters

d_token
int = 16
Width of each feature token; must be divisible by n_heads. Real problems want 32-64.
n_blocks
int = 1
Number of (feature attention, row attention) stages; real use wants 2-6.
n_heads
int = 2
Attention heads per block.
dropout
float = 0.1
Dropout rate in the attention weights and the feed-forward network.
learning_rate
float = 1e-3
AdamW step size.
weight_decay
float = 1e-5
AdamW decoupled weight decay.
batch_size
int = 64
Mini-batch size, and the population intersample attention sees.
n_epochs
int = 60
Passes over the training set; real use wants several hundred.
early_stopping
bool = False
Hold out a validation split and restore the best weights.
validation_fraction
float = 0.15
Fraction held out when early_stopping is enabled.
patience
int = 10
Epochs without validation improvement before training stops.
categorical_features
sequence of int
Column indices holding integer-coded categorical features.
device
{"cpu", "auto", "cuda", "mps"} = "cpu"
Compute device; "auto" trades reproducibility for speed.
random_state
int = 0
Seed for initialisation, shuffling and the validation split.

Attributes

network_
torch.nn.Module
The fitted network, rebuilt lazily after unpickling.
target_mean_, target_scale_
float
Target standardisation statistics.
feature_mean_, feature_scale_
np.ndarray
Feature standardisation statistics.
loss_curve_
np.ndarray of shape (n_iter_,)
Mean training loss per epoch.
n_iter_
int
Epochs actually run.
n_features_in_
int
Number of features seen during fit.

Notes

Requires PyTorch. Install with pip install 'tuiml[torch]'.

Predictions depend on the batch, because intersample attention is not a per-row function. See the classifier's notes.

Complexity. O(n m^2 d + n b m d) per epoch, with the second term quadratic in batch_size.

When to use. Choose SAINT over FTTransformerRegressor when neighbouring rows carry information about the target -- clustered or grouped data -- and accept the extra compute for it.

References

Somepalli2021
Somepalli, G., Goldblum, M., Schwarzschild, A., Bruss, C. B., & Goldstein, T. (2021). SAINT: Improved Neural Networks for Tabular Data via Row Attention and Contrastive Pre-Training. :doi:`10.48550/arXiv.2106.01342`

See Also

SAINTClassifier

The classification counterpart.
FTTransformerRegressor

Feature attention only.
NODERegressor

Differentiable oblivious trees.
python 
>>> from tuiml.algorithms.tabular_deep import SAINTRegressor
>>> model = SAINTRegressor(d_token=8, n_heads=2)
>>> model.d_token
8
>>> "regression" in SAINTRegressor.get_capabilities()
True

Fitting requires pip install 'tuiml[torch]'; the example below is a

no-op on an install without it:

python 
>>> import numpy as np
>>> from tuiml.utils.torch_backend import has_torch
>>> rng = np.random.default_rng(0)
>>> X = rng.normal(size=(300, 4))
>>> y = np.sin(X[:, 0]) * X[:, 1]
>>> if has_torch():
...     model = SAINTRegressor(n_epochs=300, random_state=0).fit(X, y)
...     print(float(model.score(X, y)) > 0.8)
... else:
...     print(True)
True

Methods

get_capabilities (cls) -> List[str]

Return supported capabilities.

get_complexity (cls) -> str

Return complexity analysis.

get_references (cls) -> List[str]

Return academic citations.