T-JEPA: Augmentation-Free Self-Supervised Learning for Tabular Data

## Abstract

Abstract Self-supervision is often used for pre-training to foster performance on a downstream task by constructing meaningful representations of samples. Self-supervised learning (SSL) generally involves generating different views of the same sample and thus requires data augmentations that are challenging to construct for tabular data. This constitutes one of the main challenges of self-supervision for structured data. In the present work, we propose a novel augmentation-free SSL method for tabular data. Our approach, T-JEPA, relies on a Joint Embedding Predictive Architecture (JEPA) and is akin to mask reconstruction in the latent space. It involves predicting the latent representation of one subset of features from the latent representation of a different subset within the same sample, thereby learning rich representations without augmentations. We use our method as a pre-training technique and train several deep classifiers on the obtained representation.
Our experimental results demonstrate a substantial improvement in both classification and regression tasks, outperforming models trained directly on samples in their original data space. Moreover, T-JEPA enables some methods to consistently outperform or match the performance of traditional methods likes Gradient Boosted Decision Trees. To understand why, we extensively characterize the obtained representations and show that T-JEPA effectively identifies relevant features for downstream tasks without access to the labels. Additionally, we introduce regularization tokens, a novel regularization method critical for training of JEPA-based models on structured data.

## 1 Introduction

Self-supervised learning has caught increasing attention in recent years due to its significant success in many applications. Self-supervision is often used for pre-training to improve models’ performance on downstream tasks. In short, the objective of self-supervision for representation learning is to generate meaningful representations from unlabeled data by using pseudo-label. By pushing dissimilar samples farther away while reducing the distance between samples that are alike, self-supervised learning can facilitate learning for both supervised and unsupervised tasks.

Self-supervision often involves generating different views of the same sample to construct positive and possibly negative samples. The term positive sample designates samples related to one another, e.g., two pictures of a dog or the same picture cropped differently. In contrast, negative samples include unrelated samples, e.g., a picture of a cat and a dog.
Given this terminology, two classes of self-supervised algorithms exist. Contrastive learning methods include negative and positive samples and non-contrastive learning techniques that rely exclusively on positive samples.
Both approaches have offered promising results by generating meaningful representations of data (Chen et al., ; He et al., ; Tian et al., ; Assran et al., ) that allow the improvement of several models’ performances on a broad range of tasks. Moreover, apart from improving supervised and unsupervised models’ performance, (Hendrycks et al., ) have shown that self-supervision also helps improve robustness and uncertainty estimation for anomaly detection tasks.

Most self-supervised approaches include deep models, which have excelled in applications that include images or text. However, using neural networks for tabular data still remains challenging (Shwartz-Ziv and Armon, ). Indeed, Grinsztajn et al. () discuss how the inherent heterogeneity of tabular data makes learning from this data structure using neural networks difficult. Nonetheless, recent work has investigated finding effective training procedures and novel architectures to learn from tabular data using neural networks. Recent advancement include improved training procedures (Kadra et al., ; Gorishniy et al., ; Hollmann et al., ), representation learning for tabular data (Ye et al., ; Bahri et al., ; Zhu et al., ) or novel architectures (Somepalli et al., ; Kossen et al., ). Despite these recent advancements, leveraging self-supervised learning for tabular data remains strenuous as most methods involve data augmentations to construct multiple views of the same data sample. While augmentations can be relatively straightforward for images or text data, constructing meaningful augmentation for tabular data is non-trivial. Augmentations for image samples often include cropping, rotation, or color alteration, while for text samples, this can include token masking or token replacement. These corruptions or augmentations of samples are domain-specific and hard to translate for structured tabular data as they can generate samples outside the data manifold.

Figure: Figure 1: T-JEPA training pipeline. In step (a) a sample $\mathbf{x}\in\mathbb{R}^{d}$ is pre-processed and masked as detailed in equation 1 and fed to the context encoder to obtain a representation in $\mathbb{R}^{l_{\mathbf{m}}\times h}$ where $l_{\mathbf{m}}$ is the number of unmasked features for context mask $\mathbf{m}$. In step (b) the unmasked representation of sample $\mathbf{x}$ is fed to the target encoder and the features’ representations are selected according to the corresponding target masks, as shown in equation 4. In step (c) the output of the context encoder is fed to the predictor to obtain a prediction for each target mask used in step (b). In step (d) we compute the $\ell_{2}$-distance between the target representations and their predictions.

As discussed by Assran et al. (), self-supervised learning for representation learning includes three types of approaches. First, Joint-Embedding Architecture (JEA) usually involves two encoders that learn to output similar embeddings for similar samples while ensuring distant embeddings for dissimilar samples. Second, Generative Architecture that aims at reconstructing a sample from a corrupted version of this sample, e.g., mask reconstruction. Third, Joint-Embedding Predictive Architecture (JEPA) resembles Generative Architecture as it relies on a similar task but in the latent space rather than the original data space. JEPA-based method consist in predicting a sample’s representation in the embedding space from the embedded representation of a corrupted version of this sample.
Recently, Assran et al. () have proposed I-JEPA, a novel self-supervised approach targeted for images that does not involve augmentations. Their approach used for pre-training offered significant improvement in classification tasks on images. Following their path, other works have extended their approach to video (Bardes et al., ), audio (Fei et al., ), and graphs (Skenderi et al., ). The present work aims to adapt JEPA for tabular data as a pre-training model to foster performance on classification and regression tasks. Recent work (Thimonier et al., ; Kossen et al., ) has demonstrated that mask reconstruction can be a relevant pretext task for representation learning of tabular data. This work extends this mask reconstruction paradigm from the data space to the latent space. Adapting such approach to structured data is particularly relevant as it avoids constructing ad-hoc data augmentations that are challenging to construct for this data type.

The main contributions of our work are the following:

- •
We put forward Tabular-Joint-Embedding Predictive Architecture (T-JEPA), a novel augmentation-free self-supervised method for tabular data.
- •
T-JEPA offers significant improvement in performance for classification and regression tasks for tabular data. Moreover, we show that augmented by T-JEPA some deep methods consistently outperform Gradient Boosted Decision Trees on the tested datasets.
- •
We extensively characterize the obtained representations and provide explanation as to why our approach enhances performance on supervised tasks.
- •
We empirically uncover a novel regularization method, regularization tokens, that is critical to escape collapsed training regimes.

## 2 Related Work

##### Self-Supervised Learning for Representation Learning

Representation learning consists in finding a transformation of the input data into a new feature space where relevant information is preserved or enhanced while noise and irrelevant details are filtered out or minimized. To that end, self-supervised approaches have become prevalent in the field.
In the field of computer vision, methods like SwAV (Caron et al., ), VICReg (Bardes et al., ) or
Barlow Twins (Zbontar et al., ) aim at producing two views of the same sample passed through two different networks, such that the outputs are maximally correlated. SwaV (Caron et al., ), for instance, aims at pushing the embeddings of different samples to belong to different clusters on the unit sphere. Barlow Twins (Zbontar et al., ) involve training two identical neural networks simultaneously on the same data but with different augmentations. The objective is to minimize the redundancy between the representations learned by each network while maximizing their agreement on the same input. VICReg (Bardes et al., ) encourages the model to focus on learning invariant features by explicitly modeling and minimizing the variance of feature embeddings. Other methods like MoCo (He et al., ) or SimCLR (Chen et al., ) focus on learning representations by contrasting positive and negative pairs. Other data structures have also benefited from representation learning using self-supervised approaches such as video (Jabri et al., ; Zhang and Crandall, ; Bardes et al., ), audio (Mittal et al., ; Niizumi et al., ; Korbar et al., ; Fei et al., ) or graph (Skenderi et al., ; You et al., ; Hwang et al., ).

Self-supervised methods can be categorized into three types of approaches: joint-embedding architectures, generative architectures, or joint-embedding predictive architectures. While the former two have been the most prevalent in the literature, recent work has demonstrated the potential of joint-embedding predictive architectures. Recently, I-JEPA (Assran et al., ) targeted for images has shown significant performance improvement over several self-supervised methods. Their approach was adapted to other data types such as video (Bardes et al., ), audio (Fei et al., ), and graphs (Skenderi et al., ) and proved to offer competitive performance in comparison with existing methods.

##### Representation Learning for Tabular Data

Representation learning for tabular data has caught increasing attention in recent years. Gorishniy et al. () extensively investigate the benefits of pre-training models on tabular data to enhance performance. In other work, Somepalli et al. (); Kossen et al. () propose a pretraining procedure to foster the performance of their novel transformer-based architectures. Parallel to that, some works have focused entirely on proposing self-supervised methods for representation learning of tabular data. One of the first approaches, VIME (Yoon et al., ), proposes to augment the existing reconstruction task with estimating mask vectors from corrupted tabular data. Bahri et al. () propose Scarf, a simple method based on contrastive learning in which different views of a sample are obtained by corrupting a random subset of features. Recent work has also investigated prototype-based representation learning for tabular data as PTaRL (Ye et al., ). Parallel to this, Ucar et al. () propose SubTab that divides the input features to multiple subsets and to perform a pretext task close to mask reconstruction. Other works, such as XTab (Zhu et al., ), TransTab (Wang and Sun, ) or UniTabE (Yang et al., ), propose self-supervised representation learning for cross-table pretraining. Wu et al. () discuss the concepts of salient and mutual information and emphasize their key role in producing meaningful sample representations. They propose SwitchTab, which aims to foster the decoupling between the salient and mutual information contained in a sample to produce its representations. Finally, Lee et al. () emphasize the necessity of correctly handling the heterogeneous features of tabular data. Somewhat close to our approach, their method consists in binning the values of each feature and proceeds to use as a pretext task the reconstruction of the bin indices rather than the value in the original feature space.

## 3 Method

As displayed in Figure 1, T-JEPA involves three main modules used to learn the final representation: a context encoder, a target encoder, and a prediction module. In short, we predict from a subset of features of a sample $\mathbf{x}$ the latent representation of another non-overlapping subset of features of $\mathbf{x}$. The context encoder is used for the prediction, while the target encoder is used to construct the representations to be predicted.

Formally, let $\mathbf{x}\in\mathcal{X}\subseteq\mathbb{R}^{d}$ be a sample with $d$ features, which can be either numerical or categorical. Let $h$ designate the hidden dimension of the transformer encoders, $f_{\theta}$ the context encoder, $f_{\bar{\theta}}$ the target encoder and $g_{\phi}$ the predictor.

##### Embedding Layers and Masking

Before being fed to the different modules, data is pre-processed using embedding layers. We normalize numerical features to obtain 0 mean and unit variance, while we use one-hot encoding for categorical features. At this point, each feature $j$ for $j\in\{1,...,d\}$ has an $e_{j}$-dimensional representation, $\mathbf{E}(\mathbf{x}_{j})\in\mathbb{R}^{e_{j}}$ , where $e_{j}=1$ for numerical features and for categorical features $e_{j}$ corresponds to their cardinality. Each sample is accompanied by a masking vector $\mathbf{m}\in\{0,1\}^{d}$ in which each entry designates whether a feature is masked: $\mathbf{m}^{j}=\mathbbm{1}\{\text{feature }j\text{ is masked}\}$, where $\mathbbm{1}\{\cdot\}$ is the indicator function. When masked, we drop the corresponding feature, and only keep the remaining unmasked features. For a mask $\mathbf{m}$ with $l_{\mathbf{m}}$ unmasked features, i.e. $d-\|\mathbf{m}\|_{1}=l_{\mathbf{m}}$, sample $\mathbf{x}$ has the following embedded representation

$$ $\tilde{\mathbf{E}}(\mathbf{x})=\{\mathbf{E}(\mathbf{x}_{j}):i\in\{1,...,d\},\mathbf{m}^{j}=0\}.$ (1) $$

Each of the $d$ features is equipped with a learned linear layer, $\texttt{Linear}(e_{j},h),\forall j\in\{1,\dots,d\}$, that embeds the $e_{j}$-dimensional representation into an $h$-dimensional space. We pass each of the $l_{\mathbf{m}}$ unmasked features’ encoded representations through their corresponding linear layers. We also learn $h$-dimensional index and feature-type embeddings. Both are added to the embedded representation of sample $\mathbf{x}$. Let $\mathbf{z}_{\mathbf{x}}^{\mathbf{m}}\in\mathbb{R}^{l_{\mathbf{m}}\times h}$ denote the obtained embedded representation of sample $\mathbf{x}$ with mask $\mathbf{m}$.

##### Regularizing Token

We also include a regularizing token [REG] inspired from the register token first proposed in (Darcet et al., ) for ViT’s. We append this token to the obtained $\tilde{\mathbf{E}}(\mathbf{x})$ representation displayed in equation 1. This token is also equipped with a learned embedding layer. This token is only used to train T-JEPA and is discarded when training supervised classifiers on the downstream task.
We later discuss the necessity of including such token in section 5.2 and observe that it acts as a regularizing method to escape training regimes leading to representation collapse. For simplicity, we do not explicitly include the regularizing token in the rest of the method description hereafter.

##### Masking strategy

The masking strategy differs between the context and target encoders. We sample several masks for each sample. Context masks are used to mask the samples before feeding them to the embedding module and context encoder. On the contrary, the target masks are used to construct the target representation after passing them through the embedding module and target encoder. Note that in both context and target masking, the regularizing token is never masked. Hence, the input of the context encoder is a masked representation of a sample $\mathbf{x}$, $\mathbf{z}_{\mathbf{x}}^{\mathbf{m}}$. In contrast, the target encoder receives as an input the embedded representation $\mathbf{z}_{\mathbf{x}}^{\mathbf{0}_{d}}$, where $\mathbf{0}_{d}$ is the $d$-dimensional null vector. Then, the target mask is used to mask the corresponding encoded feature representations of $\mathbf{z}_{\mathbf{x}}^{\mathbf{0}_{d}}$ as shown in equation 4. For both context and target encoders, we set a minimum and maximum share of features to be masked simultaneously and randomly sample a share in that interval. We construct context and target masks so that there is no overlap between context and target masks but allow overlap within context masks and target masks. Let $M_{\text{context}},M_{\text{target}}$ denote the set of sampled context and target masks, respectively.

##### Context and Target Encoders

The context encoder $f_{\theta}$ is a transformer encoder composed of $\ell$ layers and $k$ attentions heads. The context encoder relies on multi-head self-attention to produce meaningful representations for each sample. It receives as input a mask representation $\mathbf{z}_{\mathbf{x}}^{\mathbf{m}}\in\mathbb{R}^{l_{\mathbf{m}}\times h}$ and outputs a representation of similar dimension.
The target encoder’s architecture exactly reproduces the one of the context encoder. Let $f_{\bar{\theta}}$ denote the target encoder which receives as input $\mathbf{z}_{\mathbf{x}}^{\mathbf{0}_{d}}$ an unmasked embedded representation of sample $\mathbf{x}$. Like the context encoder, it outputs a representation of the same dimension as its input.

$$ $\displaystyle h_{\text{context}}^{\mathbf{m}}=f_{\theta}(\mathbf{z}_{\mathbf{x}}^{\mathbf{m}})\in\mathbb{R}^{l_{\mathbf{m}}\times h}$ $\displaystyle\quad(\text{context})$ (2) $\displaystyle h_{\text{target}}=f_{\bar{\theta}}(\mathbf{z}_{\mathbf{x}}^{\mathbf{0}_{d}})\in\mathbb{R}^{d\times h}$ $\displaystyle\quad(\text{target})$ (3) $$

Let $h_{\text{target}}^{\mathbf{m}_{k}}$ denote the masked target representation for mask $\mathbf{m}_{k}$, obtained by discarding the masked features’ representations from $h_{\text{target}}$ as done in equation 1,

$$ $h_{\text{target}}^{\mathbf{m}_{k}}=\{h_{\text{target}}^{(i)}:i\in\{1,...,d\},\mathbf{m}_{k}^{i}=0\},$ (4) $$

where $h_{\text{target}}^{(i)}$ is the $h$-dimensional representation of the $i$-th feature in the target vector $h_{\text{target}}$.
Following previous work (Assran et al., ), The parameters of the context encoder, $\theta$, are learned through gradient-based optimization. In contrast, the parameters of the target encoder $\bar{\theta}$ are updated via an exponential moving average (EMA) of the context encoder’s parameters.

##### Predictor

The predictor $g_{\phi}$ is also set to be a transformer encoder whose weights are conjointly learned with the context encoder’s weights through gradient-based optimization. The predictor’s hidden dimension is downsized from the encoders’ dimension $h$, using a linear layer. The predictor takes as input $h_{\text{context}}^{\mathbf{m_{j}}}$, the output of the context encoder for mask $\mathbf{m}_{j}\in M_{\text{context}}$, and a target mask $\mathbf{m}_{k}\in M_{\text{target}}$ designating which features to be predicted, $g_{\phi}(h_{\text{context}}^{\mathbf{m_{j}}},\mathbf{m}_{k})$. We parameterize the mask tokens in $\mathbf{m}_{k}$ by a learnable vector to which is added a positional embedding. Each context output is passed $|M_{\text{target}}|$ times through the predictor module to predict the corresponding feature representation for each target mask.

##### Loss

The loss function used to optimize the weights $\theta,\phi$, of the context encoder and predictor respectively, is the $\ell_{2}$-distance between the reconstructed representation, $g_{\phi}(h_{\text{context}}^{\mathbf{m}},\mathbf{m}_{k})$ and the target representation $h_{\text{target}}^{\mathbf{m}_{k}}$,

$$ $\mathcal{L}(\mathbf{x};M_{\text{context}},M_{\text{target}})=\frac{1}{|M_{\text{target}}|}\cdot\frac{1}{|M_{\text{context}}|}\sum_{\mathbf{m}\in M_{\text{context}}}\sum_{\mathbf{m}_{k}\in M_{\text{target}}}\left\|g_{\phi}(h_{\text{context}}^{\mathbf{m}},\mathbf{m}_{k})-h_{\text{target}}^{\mathbf{m}_{k}}\right\|_{2}^{2}.$ (5) $$

## 4 Experiments

### 4.1 Experimental Setting

##### Datasets

Following previous work (Ye et al., ), we experiments on $6$ datasets with heterogeneous features to test the effectiveness of T-JEPA. We test our approach on several supervised tabular deep learning tasks such as binary and multi-class classification, as well as regression. We use as performance metrics Accuracy ($\uparrow$) and RMSE ($\downarrow$) for classification and regression respectively. The datasets we include in our experiments are Adult (AD) (Kohavi et al., ), Higgs (HI) (Vanschoren et al., ), Helena (HE) (Guyon et al., ), Jannis (JA) (Guyon et al., ), ALOI (AL) (Geusebroek et al., ) and California housing (CA) (Pace and Barry, ). We summarize the characteristics of all $6$ datasets in Table 5 in appendix B.

##### Baselines

To assess whether our self-supervised approach can foster performance on tabular applications, we compare the performance of several widely used tabular approaches with and without our self-supervised pre-training. We test our method on MLP (Taud and Mas, ), DCNV2 (Wang et al., ), ResNet (He et al., ), AutoInt (Song et al., ) and FT-Transformer (Gorishniy et al., ). For completeness, we also compare the performance on the downstream task to other SSL methods like PTaRL (Ye et al., ), SwitchTab (Wu et al., ), BinRecon (Lee et al., ) and VIME (Yoon et al., ) (see Table 6 in Appendix E). Similarly, as often considered as the go-to methods for supervised tasks on tabular data, we also compare the performance of T-JEPA to XGBoost (Chen and Guestrin, ) and CatBoost (Prokhorenkova et al., ).

##### T-JEPA Training

We split each dataset into training/validation/test sets (80/10/10) which were used for selecting both the hyperparameters of T-JEPA and of the models used for the downstream task.
More precisely, in a model-agnostic manner, we relied on a systematic approach to train and evaluate the embedding space generated by T-JEPA. First, following previous work (Assran et al., ; Bardes et al., ) T-JEPA was trained on the training set and we conducted a hyperparameter tuning using a linear probe on the validation set to select the best configuration for each dataset. Second, we used the trained context encoder to generate data representations on which the subsequent model were trained. We refer the reader to appendix A.2 for more details on hyperparameter selection. Every experiment can be reproduced with the code made available on GitHub.

##### Downstream task

To adapt the T-JEPA representations to each model’s input dimensions, we added a projection layer. We experimented with several techniques, including linear flattening, linear per-feature transformation, convolutional projections, and max and mean pooling. We refer the reader to appendix A.3 for more details. The projection layer was jointly trained with the downstream task and tailored to each model. Hyperparameter tuning was performed based on validation set performance, and final evaluations were conducted on the test set, comparing results to models trained on the original dataset representations. To ensure the robustness and reliability of our findings, we conducted the experiment 20 times. We report the mean and standard deviation of the performance metrics.

**Table 1: Performance metrics for different models with and without T-JEPA representation across various datasets. Metrics for the models without T-JEPA are averaged over $10$ runs and taken from Ye et al. (). Regarding models augmented by T-JEPA, we report an average over $20$ runs, see Table 6 in appendix E for comparison with existing methods. We report in bold the metric that wins between the raw data representation and T-JEPA’s. We report in blue the best performing deep method and underline the overall best metric for a dataset.**
|  | AD $\uparrow$ | HE $\uparrow$ | JA $\uparrow$ | AL $\uparrow$ | CA $\downarrow$ | HI $\uparrow$ | Wins |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Baseline Neural Networks |  |  |  |  |  |  |  |
| MLP | $0.825$ | $0.352$ | $0.672$ | $0.917$ | $0.518$ | $\mathbf{0.681}$ | $1$ |
| +T-JEPA | $\mathbf{0.866}$ | $\mathbf{0.400}$ | $\underline{\mathbf{0.728}}$ | $\mathbf{0.961}$ | $\mathbf{0.468}$ | $0.517$ | $\mathbf{5}$ |
| DCNv2 | $0.826$ | $0.34$ | $0.662$ | $0.905$ | $0.502$ | $\mathbf{0.681}$ | $1$ |
| +T-JEPA | $\mathbf{0.861}$ | $\mathbf{0.399}$ | $\mathbf{0.723}$ | $\mathbf{0.955}$ | $\underline{\mathbf{0.420}}$ | $0.525$ | $\mathbf{5}$ |
| ResNet | $0.813$ | $0.354$ | $0.666$ | $0.919$ | $0.537$ | $0.682$ | 0 |
| +T-JEPA | $\mathbf{0.865}$ | $\underline{\mathbf{0.401}}$ | $\mathbf{0.718}$ | $\underline{\mathbf{0.964}}$ | $\mathbf{0.441}$ | $\mathbf{0.705}$ | $\mathbf{6}$ |
| AutoInt | $0.823$ | $0.338$ | $0.653$ | $0.894$ | $0.507$ | $\mathbf{0.685}$ | $1$ |
| +T-JEPA | $\mathbf{0.866}$ | $\mathbf{0.351}$ | $\mathbf{0.710}$ | $\mathbf{0.938}$ | $\mathbf{0.448}$ | $0.517$ | $\mathbf{5}$ |
| FT-Trans | $0.827$ | $0.363$ | $0.675$ | $0.913$ | $0.486$ | $\mathbf{0.689}$ | $1$ |
| +T-JEPA | $\mathbf{0.864}$ | $\mathbf{0.384}$ | $\mathbf{0.708}$ | $\mathbf{0.921}$ | $\mathbf{0.444}$ | $0.551$ | $\mathbf{5}$ |
| Gradient Boosted Decision Trees (GBDT) |  |  |  |  |  |  |  |
| XGBoost | $\underline{0.874}$ | $0.368$ | $0.720$ | $0.951$ | $0.462$ | $\underline{0.729}$ | N/A |
| CatBoost | $0.873$ | $0.381$ | $0.721$ | $0.946$ | $0.430$ | $0.726$ | N/A |

### 4.2 Results

##### Binary and Multi-Class Classification

As displayed in Table 1, for the classification tasks, T-JEPA pre-training improved the performance of all evaluated models except on the Higgs (HI) dataset where only the performance of ResNet improves when augmented by T-JEPA. On average T-JEPA increased the performance of MLP by $6.8$% excluding HI and $0.7\%$ otherwise, DCNv2 by $7.5\%$ and $1.4\%$, ResNet by $7.1\%$ and $6.4\%$, AutoInt by $5.8\%$ and $-0.32\%$ and FT-Transformer by $3.6\%$ and $-1.1\%$.
Interestingly, despite their design advantages for tabular data, recent attention-based models like FT-Transformer and AutoInt, though improved by T-JEPA, did not surpass architectures like ResNet and MLP, which exhibited the most significant gains with T-JEPA. This outcome suggests that T-JEPA may complement the inductive biases of more traditional architectures like ResNet and MLP, which, despite their simpler design, are better equipped to leverage the feature representations learned through T-JEPA. These findings underscore the importance of considering both model architecture and pre-training strategies when optimizing for tabular data tasks, highlighting T-JEPA’s potential as a tool for enhancing model performance. Finally, we observe that for classification tasks ResNet+T-JEPA appears as a strong candidate as it outperforms all other methods including GBDTs, which is in line with Gorishniy et al. () and Zabërgja et al. ().

##### Regression

In regression tasks, T-JEPA induces a consistent decrease of the RMSE on the California Housing (CA) dataset for all tested models, indicating improved performance. For the MLP model, the MSE improved by $9.7\%$, DCNv2 by $16.3\%$, ResNet by $17.9\%$, AutoInt by $11.7\%$ and FT-Transformer by $8.6\%$. This indicates that T-JEPA pre-training benefits the learning process for regression tasks.

##### Compute

T-JEPA pre-training consistently improves performance across a variety of deep learning models, demonstrating its model-agnostic benefits while also requiring moderate compute and training time as displayed in Table 5 in appendix B.

## 5 Discussion

### 5.1 Representation space evaluation

#### 5.1.1 Metrics

To assess the characteristics of the generated representation space, we rely on three metrics: Kullback-Leibler divergence, uniformity score and feature variances. These metrics provide insights into the distribution, compactness and separability of the embedded features.

##### Kullback-Leibler divergence

The Kullback-Leibler (KL) divergence measures the difference between the probability distributions of random points within the embedding space. Formally, given two probability distributions $P$ and $Q$ over the same variable $x$, the KL-divergence from $Q$ to $P$ is defined as:

$$ $D_{\text{KL}}(P\parallel Q)=\sum_{x}P(x)\log\frac{P(x)}{Q(x)}.$ (6) $$

We consider random points $x_{i}=\texttt{flatten}(\mathbf{h}_{i})$ and $x_{j}=\texttt{flatten}(\mathbf{h}_{j})$ from the embedding space and estimate their distributions $P$ and $Q$, where flatten is an operation that converts the high-dimensional embeddings $\mathbf{h}\in\mathbb{R}^{d\times h}$ into a one-dimensional vector $\mathbf{x}\in\mathbb{R}^{d\cdot h}$. A lower KL-divergence indicates that the embedded space has consistent and similar distributions for different regions. For completeness we also include the euclidean distance between the flattened representations.

##### Uniformity

We rely on the uniformity score (Wang and Isola, ) to evaluate the preservation of maximal information within the feature distribution. This score leverages the Gaussian potential kernel $G_{t}:\mathcal{S}^{d}\times\mathcal{S}^{d}\rightarrow\mathbb{R}_{+}$, defined as $G_{t}(u,v)\triangleq e^{-t\|u-v\|^{2}_{2}},\;t>0$. The uniformity score is defined as the logarithm of the average pairwise Gaussian potential, formally:

$$ $\texttt{uniformity}\triangleq-\log\underset{x,y\stackrel{{\scriptstyle\text{i.i.d.}}}{{\sim}}p_{\text{data}}}{\mathbb{E}}[G_{t}(u,v)]=-\log\underset{x,y\stackrel{{\scriptstyle\text{i.i.d.}}}{{\sim}}p_{\text{data}}}{\mathbb{E}}[e^{-t\|u-v\|^{2}_{2}}],\;t>0.$ (7) $$

This metric is intricately connected to the notion of uniform distribution on the unit hypersphere. This uniformity score allows us to obtain a nuanced measure that captures the degree of information preservation in the feature distribution.

##### Feature variance

Finally, we also include intra-feature variance as a measure of the variability within individual features, and inter-feature variance to quantify the variability of the hidden dimension across different features. More precisely:

$$ $\begin{array}[]{cc}\sigma_{\text{intra}}^{2}=\frac{1}{d}\sum_{i=1}^{d}\text{Var}(\mathbf{h}_{i,:}),&\sigma_{\text{inter}}^{2}=\frac{1}{h}\sum_{j=1}^{h}\text{Var}(\mathbf{h}_{:,j}),\end{array}$ (8) $$

where $\mathbf{h}\in\mathbb{R}^{d\times h}$ represents the embedding space with $d$ features and $h$ hidden dimensions.

#### 5.1.2 Embedding space characterization

**Table 2: Key metrics tracking T-JEPA’s representation learning on the Jannis (JA) dataset over different training epochs.**
| Metric/Epoch | 0 | 20 | 40 | 60 | 80 | 100 | 120 |
| --- | --- | --- | --- | --- | --- | --- | --- |
| $\sigma_{\text{intra}}^{2}$ | $1.00$ | $1.03$ | $1.08$ | $1.12$ | $1.14$ | $1.17$ | $1.17$ |
| $\sigma_{\text{inter}}^{2}$ | $0.65$ | $0.16$ | $0.14$ | $0.09$ | $0.08$ | $0.07$ | $0.07$ |
| $D_{KL}$ | $9.30e^{-4}$ | $3.61e^{-4}$ | $6.62e^{-2}$ | $5.06e^{-2}$ | $6.55e^{-2}$ | $7.66e^{-2}$ | $9.38e^{-2}$ |
| $\|\cdot\|_{2}$ | $5.83$ | $3.93$ | $54.3$ | $50.6$ | $68.1$ | $71.4$ | $70.0$ |
| uniformity | $3.12$ | $0.94$ | $10.69$ | $11.20$ | $11.23$ | $11.32$ | $11.38$ |

Figure: Figure 2: Visualization of the representation space at various epochs of the T-JEPA pretraining on the Jannis (JA) dataset. Each plot depicts the density of transformed points in two dimensions, with darker areas indicating higher density.
Refer to caption: images/uniform_analysis.png

To assess whether T-JEPA generates representations uniformly spanned across the embedding space, we display in Figure 2 the obtained sample representations for the JA dataset at different training epochs. We randomly sample 50,000 points and rely on the PaCMAP dimensionality reduction technique to reduce their representations in the embedding space to two dimensions for visualization purposes. We observe in Figure 2 that T-JEPA effectively utilizes the representation space, evolving from an initial collapsed distribution towards a more structured arrangement as the training progresses. This behavior is indicative of the model’s capacity to learn distinct and meaningful representations, which are crucial for downstream tasks.
We also display in Table 2 how the intra-feature and inter-feature variances ($\sigma_{\text{intra}}^{2}$, $\sigma_{\text{inter}}^{2}$), uniformity score, KL divergence ($D_{\text{KL}}$), and the euclidean distance ($\|\cdot\|_{2}$) vary during training of. At each epoch, we randomly select 50,000 sample representations from the JA dataset and compute these metrics. For both KL-Divergence and euclidean distance, we report the average pairwise KL-Divergence and euclidean distance for the 50,000 samples previously selected.
We observe a stable intra-feature variance that indicates consistent information extraction, while the decreasing inter-feature variance shows improved feature specialization. The rising KL divergence and euclidean distance demonstrate that the model effectively avoids representation collapse, leading to more meaningful and discriminative embeddings. Finally, increasing uniformity across epochs is in line with the representations displayed in Figure 2 and demonstrates better separability.

### 5.2 Representation Collapse

As a non-contrastive self-supervised learning method, T-JEPA is prone to representation collapse as discussed in previous work (Bardes et al., ; Assran et al., ). EMA combined with a stop-gradient operation has been considered to be sufficient to prevent JEPA-based methods from leading to degenerate solutions in which all samples have the same representation. However, we observed that further regularization of T-JEPA was necessary to avoid such pitfall. Specifically, appending a regularization token [REG] to both target and context representations appeared critical to escape the initial representation collapse.

Figure: Figure 3: Training regime of Joint-Embedding Predictive Architectures on tabular (JA) and image (ImageNet-1K) data. We display on the right a randomly selected sample’s representations for each critical part of the training process, where (a) describes the initial random initialization, (b) the collapsed equilibrium, (c) the regularization effect pushing the weights outside of the collapsed equilibrium, and (d) the convergence.
Refer to caption: images/training_regime.png

##### Initial Collapse

We observe that JEPA-based models undergo an initial representation collapse due to the EMA relation between the weights of the context and label encoders. Indeed, we retrained from scratch I-JEPA (Assran et al., ) on ImageNet-1K using their official implementation and hyperparameters(^1^11For computational purposes, we reduced the batch size to $16$ and kept unchanged the rest of the hyperparameters.) and observed a similar training regime as for T-JEPA.

Figure: Figure 4: Regularization token ablation. Training loss for the JA dataset across different numbers of regularization tokens [REG].
Refer to caption: images/reg_tokens.png

First, the loss collapses close to 0 after a few iterations; then, regularization starts pushing the model’s weights toward a non-trivial equilibrium. We display in Figure 3 the training regimes of both T-JEPA and I-JEPA.

##### Regularization Token

While other JEPA methods do not appear to require further regularization, we observe that appending a regularization token [REG] to the sample’s representations is instrumental in escaping training collapse regimes. As displayed in Figure 4, when training T-JEPA on the Jannis dataset without including any regularization token, the optimization process does not manage to escape the initial collapse. On the contrary, when including one or more tokens, we can escape this initial collapse and push the weights towards a non-trivial equilibrium.

### 5.3 Comparison with Gradient Boosted Decision Trees

##### Performance Comparison

Recent work (Grinsztajn et al., ; Gorishniy et al., ) discusses how neural networks tend to struggle with structured tabular data type in comparison with other non-deep methods based on gradient-boosted decision trees (GBDT). In most scenarios, approaches such as XGBoost (Chen and Guestrin, ) or CatBoost (Prokhorenkova et al., ) surpass deep learning algorithms.
We observe in Table 1 that on most datasets, methods trained on raw data are significantly outperformed by both GBDT methods. However, once augmented by T-JEPA, most methods consistently outperform GBDT or match their performance.

##### Feature Importance

Figure: Figure 5: Pairwise comparison of feature rankings using Kendall’s $\tau$ correlation on the JA dataset. Rankings are derived from XGBoost feature importance, permutation importance, T-JEPA’s embedding variance ($\sigma_{embed}$), and a random baseline.
Refer to caption: images/kendall_matrix.png

To try and understand why T-JEPA enabled some approaches to match the performance of GBDT, we investigated whether high-variance features in the latent space correlate with feature importance for downstream tasks. The embedding variance was computed as the standard deviation of each feature’s values across hidden dimensions, normalized by the dimension-wise mean.
We ranked features by their average embedding variance across all samples and compared these rankings to those generated by traditional supervised methods, including XGBoost and permutation importance, using Kendall’s $\tau$ for rank similarity.
Overall, our analysis reveals a strong correlation between high-variance features in the T-JEPA embeddings and those identified as important by supervised methods, despite T-JEPA being trained without target labels. Figure 5 provides a detailed comparison. The embedding variance ranking ($\sigma_{\text{embed}}$) shows significant correlation with XGBoost ($\tau=0.44$, with $p\text{-value}=1.73e^{-6}$). A random generated rank is provided for comparison.
These findings suggest that T-JEPA’s self-supervised framework effectively captures key features relevant to downstream tasks without supervision. The alignment between embedding variance and traditional feature importance further highlights the model’s ability to learn meaningful data representations.

## 6 Conclusion

Overall, we have proposed a novel augmentation-free self-supervised method for representation learning that has demonstrated strong performance across both classification and regression tasks on diverse datasets. We have investigated the characteristics of the obtained representations and demonstrated that our approach is relevant as it identifies pertinent features without access to the downstream task’s target. Moreover, most methods augmented by T-JEPA outperform or match the performance of GBDT, which is often considered the go-to method for supervised tasks on tabular data. In particular, aligned with previous work that demonstrated that ResNet was a solid alternative to GBDT for tabular data (Gorishniy et al., ; Zabërgja et al., ), we observe that ResNet+T-JEPA outperforms all competing methods on most datasets, including GBDT. Finally, we empirically uncovered a novel regularization method for transformers on tabular data by including a regularization token that prevents from entering collapsed training regimes.

##### Limitations and Future Work

Our method generates representations best suited for transformer-like architecture that requires to be adapted to other architectures.
Future work may include investigating other approaches to adapting JEPA-like methods for representation learning of tabular data that would be suited for other architectures than transformers. Other possible extensions of the present work might include investigating using JEPA-like reconstruction as a pretext task for self-supervised anomaly detection on tabular data. Finally, our work has emphasized the uncanny training regimes of JEPA-like methods as non-contrastive self-supervised approaches. Further investigation to provide better theoretical insight into why non-contrastive self-supervised learning works well might enable the construction of novel designs and approaches that foster performance.

##### Acknowledgement

This work was performed using HPC resources from the "Mésocentre" computing center of CentraleSupélec and École Normale Supérieure Paris-Saclay supported by CNRS and Région Île-de-France ([http://mesocentre.centralesupelec.fr/](http://mesocentre.centralesupelec.fr/)). This work was also granted access to the HPC resources of IDRIS under the allocation 2024-101424 made by GENCI. This research publication is supported by the Chair "Artificial intelligence applied to credit card fraud detection and automated trading" led by CentraleSupelec and sponsored by the LUSIS company.

## Appendix A Experimental settings

This section presents the implementation details of the project.

### A.1 Programming environment

The code environment for this project was implemented using Python and several third-party libraries. Table 3 details the main libraries used along with their respective versions. The training was done on a single NVIDIA HGX A100 GPU with 40GB of memory.

**Table 3: Main libraries used in the project.**
| Library | Description |
| --- | --- |
| Python v3.12.2 | The programming language used for the project |
| einops v0.8.0 | A flexible and powerful tool for tensor operations |
| matplotlib v3.8.4 | A library for creating static, animated, and interactive plots |
| numpy v2.1.0 | Fundamental package for scientific computing with arrays |
| pandas v2.2.2 | Data manipulation and analysis tool |
| pytorch_lightning v2.2.1 | A PyTorch wrapper for high-performance deep learning research |
| scikit_learn v1.4.1.post1 | Machine learning library for Python |
| scipy v1.14.1 | Library for scientific and technical computing |
| torch v2.3.0.post301 | PyTorch deep learning library |
| torchinfo v1.8.0 | Module to show model summaries in PyTorch |
| tqdm v4.66.2 | Progress bar utility for Python |
| xgboost v2.1.1 | Optimized gradient boosting library |

### A.2 Hyperparameter search

We employed Bayesian optimization to tune the hyperparameters of T-JEPA. The batch size was fixed at $512$ for all configurations, while the exponential moving average (EMA) decay rate was set to vary from $0.996$ to $1$. Additionally, we used four prediction masks throughout the training process. For optimization, we selected the AdamW optimizer [Loshchilov and Hutter, ] due to its proven robustness in large-scale models. The learning rate was adaptively adjusted using a cosine annealing scheduler [Loshchilov and Hutter, ], which gradually reduced it from the initial value to a minimum, $\eta_{\min}=0$.

**Table 4: Hyperparameter Configuration for Bayesian Optimization**
| Parameter | Values |
| --- | --- |
| model_num_heads | $[2,4,8]$ |
| model_dim_hidden | $[2,4,8,16,32,64,128]$ |
| model_num_layers | $[1,2,3,4,5,6,7,8,16]$ |
| model_dim_feedforward | $[64,128,256,512,768,1024]$ |
| model_dropout_prob | $(0.0,0.01)$ |
| exp_lr | $(0.00001,0.001)$ |
| mask_min_ctx_share | $(0.07,0.15)$ |
| mask_max_ctx_share | $(0.2,0.9)$ |
| mask_min_trgt_share | $(0.05,0.20)$ |
| mask_max_trgt_share | $(0.2,0.9)$ |
| pred_num_layers | $[2,4,8,16,24,32]$ |
| pred_embed_dim | $[4,8,16,32,64,128]$ |
| pred_num_heads | $[2,4,8]$ |
| pred_p_dropout | $(0.0,0.01)$ |

The hyperparameters displayed in table 4 correspond to the following:

- •
model_num_heads: number of attention heads of the context encoder $f_{\theta}$.
- •
model_dim_hidden: hidden dimension of the context encoder $f_{\theta}$.
- •
model_num_layers: number of layers of the context encoder $f_{\theta}$.
- •
model_dim_feedforward: dimension of FFN in the context encoder $f_{\theta}$.
- •
model_dropout_prob: dropout probability of the context encoder $f_{\theta}$.
- •
exp_lr: learning rate.
- •
mask_min_ctx_share: Minimum share of masked feature for the context representation (see (a) in Figure 1).
- •
mask_max_ctx_share: Maximum share of masked feature for the context representation (see (a) in Figure 1).
- •
mask_min_trgt_share: Minimum share of masked feature for the target representation (see (b) in Figure 1).
- •
mask_max_trgt_share: Maximum share of masked feature for the target representation (see (b) in Figure 1).
- •
pred_num_heads: number of attention heads of the predictor $g_{\phi}$.
- •
pred_num_layers: number of layers of the predictor $g_{\phi}$.
- •
pred_embed_dim: hidden dimension of the predictor $g_{\phi}$.
- •
pred_p_dropout: dropout probability of the predictor $g_{\phi}$.

### A.3 Projection layer

The projection layer adapts the T-JEPA representation space $h\in\mathbb{R}^{d\times h}$ to the input dimensions required by the downstream models. Several projection techniques were implemented, as described below:

- •
Linear Flatten: When using the linear flatten projection, the input is flattened into a single vector $\mathbf{h}_{\text{flatten}}\in\mathbb{R}^{d\cdot h\times 1}$ and transformed through a linear projection, $\mathbf{h}_{\text{proj}}=\mathbf{W}\cdot\mathbf{h}_{\text{flatten}}+\mathbf{b}\in\mathbb{R}^{h_{\text{new}}\times 1}$, with $\mathbf{W}\in\mathbb{R}^{h_{\text{new}}\times d\cdot h}$ the weight matrix and $\mathbf{b}\in\mathbb{R}^{h_{\text{new}}\times 1}$ the bias.
- •
Linear Per-Feature: Each feature is transformed independently by applying a linear projection to each feature vector $\mathbf{h}_{i}\in\mathbb{R}^{d\times 1}$, $\mathbf{h}_{i,\text{proj}}=\mathbf{W}_{i}\cdot\mathbf{h}_{i}+\mathbf{b}_{i}\in\mathbb{R}$, where $\mathbf{W}_{i}\in\mathbb{R}^{1\times d}$ is the weight matrix and $\mathbf{b}_{i}\in\mathbb{R}$ is the bias for each feature $i$.
- •
Convolutional Projection: The convolutional encoder applies two stages of convolution followed by max pooling. For an input representation $\mathbf{x}\in\mathbb{R}^{d\times h}$, where $d$ is the number of feature and $h$ is the hidden dimension, the convolution operation is represented as: $\mathbf{x}^{\prime}=\sigma\left(\texttt{BatchNorm}\left(\texttt{Conv2D}(\mathbf{X},\mathbf{k}_{1})\right)\right)$, where $\mathbf{k}_{1}$ is a convolutional kernel, and $\sigma$ is the activation function (ReLU). After a second convolution and pooling step, the final representation is flattened into a vector and projected to the target embedding dimension.
- •
Max Pooling: The max pooling operation selects the maximum value for each feature vector $\mathbf{h}_{i}\in\mathbb{R}^{h}$: $\mathbf{h}_{i,\text{max}}=\max(\mathbf{h}_{i,j}\mid j\in[d])$.
- •
Mean Pooling: The mean pooling operation averages values for each feature vector $\mathbf{h}_{i}\in\mathbb{R}^{h}$: $\mathbf{h}_{i,\text{mean}}=\frac{1}{h}\sum_{j\in[h]}\mathbf{h}_{i,j}.$

Each projection layer is trained jointly with the downstream task and is selected based on the structure of the input data and model requirements.

## Appendix B Compute and training time

The training process was executed on a single NVIDIA A100 GPU. Despite the varying dataset sizes, we aimed to optimize compute efficiency by carefully tuning the batch size and learning rate for each experiment. As detailed in table 5, the pretraining duration across datasets ranged from 0.34 GPU-hours to 3.64 GPU-hours. These variations largely reflect the complexity of the datasets in terms of both sample size and feature dimensions. In total, the computational demand remained within acceptable limits, allowing us to complete multiple runs with reasonable turnaround times, while also maintaining a balance between model performance and resource usage.

**Table 5: Dataset characteristics and pretraining GPU-hours.**
|  | AD | HI | HE | JA | AL | CA |
| --- | --- | --- | --- | --- | --- | --- |
| Samples | 48,842 | 98,050 | 65,196 | 83,733 | 108,000 | 20,640 |
| Numerical, Categorical | $6$, $8$ | $28$, 0 | $27$, 0 | $54$, 0 | $128$, 0 | $8$, 0 |
| Classes | $2$ | $2$ | $100$ | $4$ | $1,000$ | N/A |
| Metric | Accuracy | Accuracy | Accuracy | Accuracy | Accuracy | RMSE |
| Pretraining GPU-hours | $0.84$ | $1.80$ | $1.03$ | $1.27$ | $3.64$ | $0.34$ |

## Appendix C Alternative Strategies

Other settings have been considered before the one discussed in section 3, and were discarded because they led to collapsed regimes. We considered two alternatives. First, we considered a pipeline where only the masking strategy differs, as detailed in section C.1. Second, we also considered an alternative where the masking strategy is the one detailed in section C.1, and we also modify the predictor architecture as detailed in section C.2.

### C.1 Masking Strategy

##### Masking

Following previous work on feature masking for tabular data, we considered handling masked features by keeping the sample’s representation’s dimension constant. Features of a samples are normalized similarly as detailed in section 3 such that $\mathbf{E}(\mathbf{x}_{j})\in\mathbb{R}^{e_{j}}$, where $e_{j}=1$ for numerical features and for categorical features $e_{j}$ corresponds to their cardinality. Each sample is accompanied by a masking vector $\mathbf{m}\in\{0,1\}^{d}$ in which each entry designates whether a feature is masked: $m_{j}=\mathbbm{1}\{\text{feature }j\text{ is masked}\}$. When masked, we replace the corresponding feature value with 0 and concatenate each feature representation with the corresponding mask indicator function. Hence, each feature $j$ has an $(e_{j}+1)$-dimensional representation

$$ $\tilde{\mathbf{E}}(\mathbf{x})=((\mathbf{1}_{d}-\mathbf{m})\odot\mathbf{E}(\mathbf{x}),\mathbf{m})\in\mathbb{R}^{d\times(e_{j}+1)},$ (9) $$

where $\odot$ designates the Hadamard product and $\mathbf{1}_{d}$ the $d$-dimensional unit vector.

We pass each of the $d$ features encoded representations of sample $\mathbf{x}$ through $d$ learned linear layers $\texttt{Linear}(e_{j}+1,h)$. We also learn $h$-dimensional index and feature-type embeddings. Both are added to the embedded representation of sample $\mathbf{x}$. Let $\mathbf{z}_{\mathbf{x}}^{\mathbf{m}}\in\mathbb{R}^{d\times h}$ denote the obtained embedded representation of sample $\mathbf{x}$ with mask $\mathbf{m}$.

##### Context and Target Encoders

Given this modification, the obtained context representation’s dimension differ from the one given in equation (2).

$$ $h_{\text{context}}^{\mathbf{m}}=f_{\theta}(\mathbf{z}_{\mathbf{x}}^{\mathbf{m}})\in\mathbb{R}^{d\times h}\quad(\text{context})$ $$

The rest of the pipeline is identical to the one given in section 3. We also considered a second alternative where the above pipeline is chosen but a different predictor architecture (see section C.2) replaces the transformer as detailed in section 3.

### C.2 Predictor

Given this alternative masking strategy, we also considered an MLP-based predictor for target representation prediction. The predictor consists of $d$ separate MLPs (one for each feature). Each MLP takes as input a representation of dimension $d\cdot h$ (the flattened representation output by the context encoder), and produces and output of dimension $h$. In particular, each MLP corresponds to one feature in particular and produces an $h$-dimensional prediction for the corresponding feature given a flattened context representation of dimension $d\cdot h$.

Let us denote $g_{\phi}=\{\text{MLP}_{i}\}_{i=1}^{d}$, a target mask $\mathbf{m}_{tgt}$ and $z_{flatten}=\texttt{flatten}(h_{context}^{\mathbf{m}})$. Then the prediction for the target mask $\mathbf{m}_{target}$ is given by,

$$ $\hat{h}_{target}^{\mathbf{m}_{target}}=\{MLP_{j}(z_{flatten}):\mathbf{m}_{target}^{j}=0\}$ (10) $$

## Appendix D Embedding feature variance

The idea behind calculating the variance of embedding features is based on the assumption that features with high variability across the embedding space are more expressive and likely to capture the underlying structure of the data. As presented in Figure 6, some features (rows in the heatmap) present more perturbations across the hidden dimensions (columns in the heatmap).

Figure: Figure 6: Embedding Variance
Refer to caption: images/embedding_variance.png

To measure this variance, let us define the embedding variance score. Let $\mathbf{x}\in\mathbb{R}^{d\times h}$ represent a point the latent space, where $d$ is the number of features and $h$ is the hidden dimension. For each hidden dimension $j$, we compute the mean $\mu_{j}$ across the $d$ features, i.e., $\mu_{j}=\frac{1}{d}\sum_{i=1}^{d}x_{i,j}$. For each feature $i$, we calculate the embedding variance score, defined as $\sigma_{i,\text{embed}}=\text{Var}(\mathbf{x_{i}}-\mu)$, where $\mathbf{x_{i}}$ is the vector representing the $i$-th-feature, and $\mu=[\mu_{0},\ldots\mu_{h}]^{T}$. Features with higher embedding variance are those that stand out more in the representation space, suggesting they may carry more relevant information.

## Appendix E Detailed Performance

Table: Table 6: Performance metrics for different models and T-JEPA enhancement across various datasets including the standard deviations. The metric for models without T-JEPA and augmented by PTaRL are taken from [Ye et al., ] in which no standard deviations were reported. We also include for comparison SSL methods like SwitchTab [Wu et al., ], BinRecon [Lee et al., ] or VIME [Yoon et al., ]. We report in bold the metric that wins between the raw data representation and the augmented representations. We report in blue the best performing deep method (including SSL methods) and underline the overall best metric for a dataset.

In Table 6, we present performance metrics for several models, including neural networks (MLP, DCNv2, ResNet, AutoInt, FT-Trans), self-supervised learning (SSL) methods, and gradient-boosted decision trees (GBDT). Regarding T-JEPA, we include standard deviations from 20 independent runs to ensure statistical robustness. Performance metrics for PTaRL [Ye et al., ], SwitchTab [Wu et al., ], and BinRecon [Lee et al., ] were primarily extracted from the original papers, ensuring similar experimental setups for consistency. For VIME, we used the official code made available online by the authors to run the experiments for HE, JA, AL, and HI datasets and report the metrics in their paper for AD and CA.

T-JEPA and PTaRL offer distinct approaches to SSL. T-JEPA focuses on learning robust joint embeddings by predicting future token states, thus enhancing latent space with high-level patterns. It consistently improves deep model performance (e.g., MLP, DCNv2, ResNet), especially in accuracy-based metrics like HE and JA. In contrast, PTaRL learns discrete prototypes from tabular data, capturing key distributional patterns that augment downstream models. PTaRL excels in settings with limited data but lacks consistent statistical variance reporting, limiting direct comparison with T-JEPA. Finally, while PTaRL displays interesting performance across datasets, their approach suffers from a significant computational complexity compared to ours.

We also compare T-JEPA to other SSL methods, such as SwitchTab [Wu et al., ], BinRecon [Lee et al., ], and VIME [Yoon et al., ]. Overall we observe that most approaches augmented by T-JEPA consistently outperform competing SSL approaches apart from PTaRL on some datasets. Nevertheless, we observe that T-JEPA appears to struggle on the HI dataset compared to the existing SSL method, which tend to perform better.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. What is $h$-dimensional index and feature-type embeddings, defined on Line 195? 
2. I think Figure 1 does not quite match the description in Section 3. In particular, in Figure 1, I get a sense that there is a single context mask for each sample, while on Line 216 and Line 241, there appears to be multiple context mask for a single sample. In general, how are these masks drawn and how do you ensure that the context / target masks don't overlap? 
3. For Equation(5), I think there is a missing superscript $j$ as $\mathbf{m_j} \in M_{context}$. 
4. I don't quite follow why we can't do masking for target features before encoding, the same as we do for context. Specifically, this creates data leak, and I suspect that the mode collapse that you demonstrate in Figure 3 is related exactly to this. For a simplified example of why the data leak happens, consider that we have context and target masks, $\mathbf{m}_1$ and $\mathbf{m}_2$ respectively. Further, for simplicity, suppose $\|\mathbf{m}_1\| = |\mathbf{m}_2|$. Then I claim that picking $f\_\theta$ to be the identity map and $f\_{\bar\theta}$ to be the appropriately constructed permutation, we can achieve 0 reconstruction error. For the context features, we have the following transformations: $\mathbf{z}\_{\mathbf{x}} \xrightarrow{\mathbf{m}_1}  \mathbf{z}\_{\mathbf{x}}^{\mathbf{m}\_1} \xrightarrow{f\_\theta}  \mathbf{z}\_{\mathbf{x}}^{\mathbf{m}\_1} $. For the target, features we have: $\mathbf{z}\_{\mathbf{x}} \xrightarrow{f\_{\bar\theta}} \[\mathbf{z}\_{\mathbf{x}, \pi(1)}, \mathbf{z}\_{\mathbf{x}, \pi(2)}, \dots, \mathbf{z}\_{\mathbf{x}, \pi(d)}\] \xrightarrow{\mathbf{m}_2}  \mathbf{z}\_{\mathbf{x}}^{\mathbf{m}\_1}$, where $\pi$ is the constructed permutation. My conjecture is that the models initially learn this, but then are thrown off by the additional regularization. However, overall, I am not sure if the formulation is correct.
soundness: 3
strengths: The authors' idea of applying JEPA for representation learning on tabular data is original. The work is also well-supported by extensive experiments. In general, coming up with a masking strategy to make JEPA work for tabular data constitutes a strong engineering contribution. Tabular data, unlike the other data modalities such as vision, speech, or text, is highly heterogeneous and complex: each feature has its own distribution, some might be categorical while other continuous with varying degrees of non-Gaussianity. Therefore, representation learning on tabular data poses a significant challenge. We have seen a number of works in the recent years in this area. The present work appears to be an important step forward: it finds that JEPA works for tabular data with little modifications, which is promising.
summary: The authors propose a self-supervised representation learning approach for tabular data based on the Joint Embedding Prediction Augmentation (JEPA) framework. Methodologically, the model is based on the work of Assran et al. (2023): inputs are masked, passed to 2 encoders, and a prediction model is trying to reconstruct one of the encoded inputs from the other. Empirically, the authors show that using the features from their model, T-JEPA, increases performance on the downstream tasks. The authors also offer some discussion, highlighting that the identified features are indeed meaningful. 

[1] Assran, Mahmoud, Quentin Duval, Ishan Misra, Piotr Bojanowski, Pascal Vincent, Michael Rabbat, Yann LeCun, and Nicolas Ballas. "Self-supervised learning from images with a joint-embedding predictive architecture." In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 15619-15629. 2023.
weaknesses: I am not convinced about the contributions of Sections 5.1 and 5.3. These sections look artificial and drawn-out to me. 

## Comments on Section 5.1
The proposed metrics in Section 5.1 have little to do with the optimized objective, Equation 5. Therefore, I have no intuition of what we want these metrics to be: what is their range? what values are good / bad? how fast / slow do we want them to converge? why do we think optimizing our objective promotes this? The authors' arguments such as
> Line 366 "A lower KL-divergence indicates that the embedded space has consistent and similar distributions for different regions."

> Line 377 "This uniformity score allows us to obtain a nuanced measure that captures the degree of information preservation in the feature distribution."

> Line 422 " stable intra-feature variance that indicates consistent information extraction, while the decreasing inter-feature variance shows improved feature specialization"

are imprecise and hand-wavy. If anything, these arguments confuse me more than provide any explanation. 

## Comments on Section 5.3
I am not sure I fully grasp what the author do in Section 5.3. As far I as I understand, they want to compare the feature importance produced by XGBoost (and permutation test) to the feature importance produced by their model. However, I don't think the quantity that the authors claim to be feature importance for their method, which is the variance of the hidden features, is a meaningful proxy. The main problem with this metric is the following. Let's say we have feature vector, $\[x_1, x_2\] \in \mathbb{R}^2$, which we encode to get $\[\mathbf{h}_1^T, \mathbf{h}_2^T\]\in\mathbb{R^{2\times h}}$. The authors propose to compute variance of $\mathbf{h}_i$ across hidden dimension and use this quantity to rank $x_i$. However, in general, $\mathbf{h}_i$ does not correspond to $x_i$. Specifically, in the transformer model, when we compute the self-attention, each of $\mathbf{h}_i$ becomes a weighted some of all the other tokens. In the very simple case, consider what happens, with hidden dimension $h=1$ and if the learned encoding is a simple permutation. Clearly, this is not correct.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. Could you provide a figure illustrating the full T-JEPA framework? It would be helpful to show how masking, context encoding, target encoding, and prediction are applied to a sample of tabular data. Simply describing the pipeline may cause confusion for some readers. For example, consider referencing the structure used in Figure 3 of the I-JEPA paper [1].
2. In Figure 4 (a) to (d), the explaination regarding how the representation figure is plotted is needed. Additionally, the paper should include more discussion on the insights derived from these plots. Specifically, what does a more diverse heatmap indicate in your experiments? I found some explanations in Appendix D, but the connection was unclear until I reached that section. At the very least, it would be helpful if the main text refers readers to Appendix D for a more detailed explanation and interpretation of the figure.
3. Is there any correlation between the performance improvements gained from T-JEPA and the number of parameters in the backbone models? Can stronger models with more parameters capture more context information and therefore leading to stronger representation structure and better performance?
4. I suggest the authors include the results comparison with PTaRL to table 1, as the two methods are all enhancements to table learning paradigm.

[1] Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture. CVPR 2023.
soundness: 3
strengths: 1. The introduction of a regularization token is an interesting idea. The authors conduct thorough ablation studies to validate its effectiveness. I appreciate that this technique is not only technically sound but also produces convincing experimental results.
2. The experimental design is comprehensive. The paper evaluates its approach across multiple dimensions: (1) testing on various backbones with and without T-JEPA, (2) comparing backbone+T-JEPA with backbone+PTaRL, (3) benchmarking against other self-supervised models, and (4) contrasting results with supervised decision tree methods. Additionally, the analysis of representation visualization and collapse is thorough.
3. The inclusion of failed framework design attempts in the appendix is particularly valuable. This transparency provides deeper insights and can guide future researchers who are to follow the work.
summary: This paper introduces T-JEPA, a novel self-supervised learning method for tabular data that does not require data augmentations. T-JEPA predicts the latent representation of one subset of features from another subset within the same sample. It operates in the latent space rather than the original data space, avoiding the need for explicit data augmentations. The authors conducts extensive experiments to demonstrate that T-JEPA can serve as a model-agnostic performance boost for tabular learning tasks. The paper also analyzes the reasons behind performance gain through the lens of representation space visualization.
weaknesses: 1. The paper lacks detailed information about the model architectures used in the main experiments. For example, does ResNet refer to ResNet-18 or ResNet-50? Providing the exact specifications of the hyperparameters can enhance the reproducibility of the experiments.
2. When compared to PTaRL, which is also enhancement for tabular learning, the performance gain is not competitive. Although the idea is novel, the lackluster performance severely limits the contribution of this paper. The authors are encouraged to reimplement PTaRL themselves, rather than simply reusing numbers from previous papers. Though PTaRL has reproducibility challenges due to it not being open-sourced, relying solely on existing published results without additional validation weakens the paper.
3. The description of the methods, specifically how masking is done, is confusing. It would be better to include a figure showcasing the framework of T-JEPA and how downstream task is performed.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: 1. What is the meaning of sentence 'We also learn h-dimensional index and feature-type embeddings' in Line 194-195? Please make it more clear.  
2. For Regularization token ablation in Fig.4, it seems that without including any regularization token (blue) has the best performance. However, the authors indicate that without including any regularization token, the optimization process does not manage to escape the initial collapse, thereby resulting in a poor performance. Please give more explanations.  
3. Except for regularization token ablation on training loss, the authors should further give ablation study on performance, since (1) training loss is not identical to the final evaluation performance and (2) the regularization token is critical for training JEPA-based models on structured data.  
4. The sentence 'We construct context and target masks so that there is no overlap between context and target masks but allow overlap within context masks and target masks' in Line 215-216 is confusing. 'no overlap between context and target masks' and 'allow overlap within context masks and target masks' are contradictory.  
5. The description about right four subfigures of Fig.3 is not enough.
soundness: 3
strengths: 1. The proposed SSL method is model-agnostic and could be coupled with various deep tabular models, which has wide range application.  
2. The strategy for masking representations is novel in tabular domain.  
3. The experiment is extensive.
summary: This paper proposes a novel augmentation-free SSL method for tabular data in a model-agnostic way, that relies on a Joint Embedding Predictive Architecture (JEPA) and is akin to mask reconstruction in the latent space.
weaknesses: Please refer to the questions below.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: - The authors emphasize the term "augmentation-free," but I wonder if previous methods such as SubTab and SCARF are also augmentation-free in the same way defined in this paper.

- There is a workshop paper, *STab: Self-supervised Learning for Tabular Data,* which also claims to be augmentation-free, and it appears they genuinely lack any form of augmentation. How does the current work’s use of “augmentation-free” compare with such approaches? Additionally, a discussion of the differences between these works in this context would be helpful for clarity.

- What are the variances in the reported results?

- It is essential to include comparisons with other SSL methods for tabular data in the main paper.

- Given recent advancements in LLM-based methods and foundation models, it might also be relevant to consider these approaches for this problem.


-----------------------------------------------------------------------
**After Rebuttal:** 
Thank you to the authors for conducting new experiments and significantly improving the experimental section compared to the submitted draft. Based on your response, my understanding of the proposed method and its comparison with the baselines was indeed correct. I have increased my score from 3 to 5, reflecting the substantial improvements made during the rebuttal process.

Unfortunately, I am unable to provide further comments on the paper. However, if the authors are able to respond, it would be greatly appreciated:

1- The authors mentioned in their response that "Results from SwitchTab and VIME were directly obtained from their paper." However, there is a mismatch between the reported values in the response section and those in the VIME paper for MNIST. Could you please elaborate on this discrepancy?

2- Could you also provide details about the capacity of each model architecture? The authors stated that in some cases, results were taken directly from the paper, and in others, the available codes were used to run experiments. It is unclear what different architectures were used for each method.
soundness: 2
strengths: The method is simple and straightforward, addressing an important question.
summary: The authors propose a modification of current SSL methods for tabular settings, such as SubTab, and demonstrate the effectiveness of their method on simple baselines; however, they did not conduct a thorough comparison with these baselines.
weaknesses: - The experiments do not seem to include some relevant baselines, including SubTab, which appears closely related to this methodology, as well as SCARF and others.

- I see that the authors are using the same dataset as the one used in prior work. However, is there a reason for not including datasets used in VIME, SCARF, and SubTab?

- The concept appears very similar to SubTab, at least from a high-level perspective, but the authors did not thoroughly discuss the distinctions. Additionally, the experiments do not demonstrate how the proposed method improves upon SubTab.

- It is unclear how much the method outperforms PTaRL, as the latter shows superior results in several cases. The authors also did not provide sufficient explanation for this, instead placing the results in the appendix.

Meta-review: This paper focuses on using self-supervised learning to process tabular data. It proposes a simple but effective method to learning deep representations in an unsupervised manner, without complex augmentation. Experiments are overall comprehensive and well justify the paper's claims. 

The strengths of this paper are reflected in several aspects. First, the research problem is interesting and important, but less explored. Second, the technical details of the proposed method are easy to follow. Third, experiments are extensive, especially after rebuttal. The main weakness of this paper is that its technical contribution is somewhat limited. AC checks the paper, review comments, and author feedback, and organizes the discussion with reviewers. Overall, this is a good paper that contributes to the research field of self-supervised learning and would inspire lots of downstream applications. Therefore, it is recommended for acceptance. The author's responses and promises in the rebuttal should be reflected in the final version.

Decision: Accept (Poster)