T-JEPA: Augmentation-Free Self-Supervised Learning for Tabular Data
Authors: Hugo Thimonier
José Lucas De Melo Costa
Fabrice Popineau
Arpad Rimmel
Bich-Liên DOAN
Venue: ICLR 2025 Poster (2025)
Keywords: Self-Supervised Learning
Tabular Data
Representation Learning
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: This work presents a novel augmentation-free self-supervised method for tabular data based on the joint-embedding predictive architecture framework.

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.

Published as a conference paper at ICLR 2025
T-JEPA: AUGMENTATION-FRER SELF-SUPERVISED
LEARNING FOR TABULAR DATA
Hugo Thimonier José Lucas De Melo Costa
Fabrice Popineau' Arpad Rimmel' Bich-Lièn Doan' Université Paris-Saclay, CNRS, CentraleSupêlec, Laboratoire Interdisciplinaire des Sciences du Numérique, 91190, Gif-sur- -ur-ste-Yvette,, France.
2 Emobot, France. (name) ) ( (surname) @cent l lesupelee lcc. fr
ABSTRACT
Self-supervision is often used for pre-training to foster performance on down- stream 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 superviinnn for structured data. In the present work, we propose a novel augmentation-free SSL method for tabular data. Our approach, T-JEPA, relies Architecture (JEPA) and is akin to mask reconstruction in the latent space. It involves predicting the latent representation of 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 out- perform 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 regular- ization tokens, novel regularization method critical for training of EEPA-based models on structured data.
Joint Embedding Predictive
INTRODUCTION
Self-supervised learning has caught increasing attention in recent years due to its significant success in many applications. Self-supervision often used for pre- 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 alike, self-supervised learning can facilitate learning for both supervised and unsupervised tasks.
Self-supervision often involves generating different views of the same sample construct positive and possibly negative samples. The term positive sample designates samples related to one another, e.g.. two pictures of dog or the same picture cropped differently. In contrast, negative samples include unrelated samples, e.g., a picture of a cat and 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., 2020; He et al., 2020; Tian et al., 2019; Assran et al., 2023) that allow the improvement of several models' performances on a broad range of tasks. Moreover, apart from improving supervised
"Equal contribution

---

Published as a conference paper at ICLR 2025
Targets pred.
Context encoder
Context
Predictor
fe
U
Original sample
(b)
(d)
S8
E2
Target encoder
Target masks
Targets rep.
Figure 1: T-JEPA training pipeline. In step (a) a sample R is pre- rre-processe and masked as detailed equation I and fed to the context encoder to obtain a representation R'm where 4m is the number of unmasked features for context mask m. In step (b) the unmasked representation of sample is fed the target encoder and the features' representations are selected according to the corresponding target masks, as shown in equation 4 and Figure 7. In step (c) the output of the context encoder fed to the predictor to obtain a prediction for each target mask used in step (b). In step (d) compute the éy-distance between the target representations and their predictions.
and unsupervised models' performance, (Hendryeks al., 2019) 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, 2021). Indeed, Grinsztajn et al. (2022) discuss how the inherent het- erogeneity of tabular data makes learning from this data structure using neural networks difficult. Nevertheless, recent work has investigated finding effective training procedures and novel architec- tures to learn from tabular data using neural networks. Recent advances include improved training procedures (Kadra et al., 2021; Gorishniy et al., 2021; Hollmann et al., 2023), representation learning for tabular data (Ye et al., 2024; Bahri et al., 2022; Zhu et al., 2023) or novel architectures (Somepalli et al., 2021; ; Kossen et al., 2021). Despite these recent advances, leveraging self-supervised learning for tabular data remains strenuous most methods involve data augmentations construct multiple views of the same data sample. While augmentations can be relatively straightforward for images 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 token replacement. These corruptions augmentations samples domain-specifi and hard to translate for structured tabular data as they can generate samples outside the data manifold.
As discussed by Assran et al. (2023), self-supervised learning for representation learning includes three types of approaches. First, Joint- Seitt-Emeedding 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 corrupted version of this sample, e.g., mask reconstruction. Third, EmbeddEm Predictive Architecture (JEPA) resembles Generative Architecture as it relies on a similar task but 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 corrupted version of this sample. Recently, Assran al. (2023) have proposed I-JEPA, novel self-supervised approach targeted for images that does not involve augmentations. Their approach used for pre-traie pre-training offered significant improvement in classification tasks images. Following their path, other works have
2

---

Published as a conference paper at ICLR 2025
extended their approach to video (Bardes et al., 2024), audio (Fei et al., 2024), and graphs (Skenderi et al., 2023). The present work aims to adapt JEPA for tabular data as a pre-training model to foster performance classification and regression tasks. Recent work (Kossen et al., 2021; Ucar et al., 2021; Thimonier et al., 2024) 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 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--
Predictive Architecture (T-JEPA), novel augmentattinn-free self-supervised method for tabular data.
T-JEPA offers significant improvement in performance for classification and regression tasks for tabular data. Moreover, 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 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 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., 2020), VICReg (Bardes et al., 2022) or Barlow Twins (Zbontar et al., 2021) aim at producing two views the same sample passed through two different networks, such that the outputs are maximally correlated. SwaV (Caron et al., 2020), for instance, aims at pushing the embeddings of different samples to belong to different clusters on the unit sphere. Barlow Twins (Zbontar et al., 2021) 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 the same input. VICReg (Bardes et al., 2022) 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., 2020) or SimCLR (Chen et al., 2020) focus on learning representations by contrasting positive and negative pairs. Other data structures have also benefited from representation learning using self- aaf-suueeeieed approaches such as video (Jabri et al., 2020; Zhang and Crandall, 2022; Bardes et al., 2024), audio (Mittal al., 2022; Niizumi et al., 2021; Korbar et al., 2018; Fei et al., 2024) or graph (Skenderi et al., 2023; You et al., 2020; Hwang et al., 2020).
Self self-super methods can be categorized into three types of approaches: joint-embedding archi- tectures, 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., 2023) targeted for images has shown significant performance improvement over several self-supervised methods. Their ap- proach was adapted to other data types such as video (Bardes al., 2024), audio (Fei al., 2024), and graphs (Skenderi et al., 2023) 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 recent years. Gorishniy et al. (2021) extensively investigate the benefits of pre-training models on tabular data to enhance performance. In other works, Somepalli et al. (2021) and Kossen et al. (2021) propose a pretraining procedure to foster the performance of their novel transformer- based architectures for tabular data. Parallel to that, some works have focused entirely on proposing sulf-super methods for representation learning of tabular data. One of the first approaches, VIME (Yoon et al., 2020), proposes augment the existing reconstruction task with
3

---

Published as a conference paper at ICLR 2025
estimating mask vectors from corrupted tabular data. Bahri et al. (2022) propose SCARF, simple method based on contrastive learning in which different views of a sample are obtained by corrupting random subset of features. Recent work has also investigated prototype- based representation learning for tabular data such as PTaRL (Ye et al., 2024). Other works, such as X'Tab (Zhu et al., 2023), TransTabb (Wang and Sun, 2022) UniTabE (Yang al., 2024), propose self-supervised representation learning for cross-table pretraining. Wu et al. (2024) 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. Lee et al. (2024) 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. Also related to our approach, Sui et al. (2024) propose an augmentation- free method that can be applied to tabular data. Their method consists in simultaneously reconstructing multiple randomly generated data projection functions to generate a data representation. Finally, most related to our method, Ucar al. (2021) propose SubTab that divides the input features to multiple subsets to perform a pretext task close to mask reconstruction. The core difference with T-JEPA lies in the fact SubTab performs mask reconstruction in the original dataspace while T-JEPA performs this task in the embedded space. We provide in Appendix F.2 a more comprehensive description of each SSL methods relevant to the present work.
METHOD
As displayed in Figure T-JEPA involves three main modules used learn the final representation: a context encoder, a target encoder, and prediction module. In short, we predict from a subset of features of sample x the latent representation of another non-overlapping subset of features of x. The context encoder is used for the prediction, while the target encoder is used to construct the representations to be predicted.
Formally, let x E X' C Rd be sample with d features, which can be either numerical or categorical. Let h designate the hidden dimension of the transformer encoders, So the context encoder, fg the target encoder and go 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 mean and unit variance, while we use one-hot encoding for categorical features. At this point, each feature for j E {1,..., d) has an ey-dimensional representation, E(x;) E RS where ej 1 for numerical features and for categorical features Ey corresponds to their cardinality. Each sample is accompanied by a masking vector m e {0, 1 in which each entry designates whether a feature is masked: m² = 1 [feature j is masked], where 1 {-} is the indicator function. When masked, we drop the corresponding feature, and only keep the remaining unmasked features. For mask m with 4m unmasked features, i.e. d - ||m|]| = 4m sample x has the following embedded representation
Exx = (E(x)) i E (I,...,d, m = -3.
(1)
Each of the d features is equipped with a learned linear layer, Li niarer h), Vj E (1,...,dy d}. that embeds the ej ey-dimensional representation into an /i-dimensional space. We pass each of the Im unmasked features' encoded representations through their corresponding linear layers. We also learn h-dimensional index and feature-type embeddings following standard practice when leveraging transformers for tabular data. Both are added to the embedded representation of sample x. Let z E R denote the obtained embedded representation of sample x with mask m.
Regularizing Token We also include a regularizing token [RRG] inspired from the register token first proposed in (Darcet et al., 2024) for ViT's. We append this token to the obtained E(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. See Fig in Appendix H for illustration. 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
4

---

Published as a conference paper at ICLR 2025
regularizing token in the rest of the method description hereafter. See Appendix H for more detailed discussion on regularization tokens.
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 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 sample x, xx In contrast, the target encoder receives an input the embedded representation z where 0g is the d-dimensional null vector. Then, the target mask is used to mask the corresponding encoded feature representations of z as shown in equation 4. For both context and target encoders, we set minimum and maximum share of features to be masked simultaneously and randomly sample a share in that interval. Let Moontext Murget denote the set of sampled context and target masks, respectively. We construct Moontest, Marget such that intra-overlaps are permitted (masks from the same set can overlap), but inter-overlaps are not permitted (a mask from Mcontext cannot overlap with mask from Marger).
Context and Target Encoders The context encoder fo is transformer encoder composed of & layers and i attentions heads. The context encoder relies on multi-head self-attention to produce meaningful representations for each sample. It receives as input mask representation z E RR' X/ and outputs a representation of similar dimension. The target encoder's architecture exactly reproduces the one of the context encoder. Let fa denote the target encoder which receives input z an unmasked embedded representation of sample x. Like the context encoder, it outputs a representation of the same dimension its input.
Im context fo(zz) E R/u (context)
Nurget = (2( dd (target)
Let Minger denote the masked target representation for mask m. obtained by discarding the masked features' representations from hargett as done equation 1,
J Iurget
[Munget i E (1, d). m 0},
where bunngt is the ht-dimensional representation of the i-th feature in the target vector Atrgett Fol- lowing previous work (Assran et al., 2023). The parameters of the context encoder, are learned through gradient-basee based optimization. In contrast, the parameters of the target encoder 0 are updated via an exponential moving average (EMA) of the context encoder's parameters.
Predictor The predictor I4 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 Accomtoos the output of the context encoder for mask mj E Moontext and a target mask m; E Murget designating which features to be predicted, go (hannteu m2). We parameterize the mask tokens in m, by learnable vector to which is added positional embedding. Each context output is passed |MMarge times through the predictor module to predict the corresponding feature representation for each target mask.
Loss The loss function used to optimize the weights 0, Φ, of the context encoder and predictor respectively, is the fy-distance between the reconstructed representation, Io (h m2) and the target representation /eagget-
L(x; M Meontet: Mearget)
1
Σ
Σ
(/ ma)
huugee (5)
Marger
M Moontext
me Monia my Marpi
EXPERIMENTS
4.1 EXPERIMENTAL SETTING
Following previous work (Ye et al., 2024), we experiments on 7 datasets with hetero- geneous features to test the effectiveness of T-JEPA. We test our approach several supervised
Datasets
5

---

Published as a conference paper at ICLR 2025
Table 1: Performance metrics for different downstream models trained on the original dataspace and the generated T-JEPA representation across datasets. We also include for completeness the performance of XGBoost (Chen and Guestrin, 2016) and CatBoost (Prokhorenkova al., 2018) as a baseline. We report an average over 20 runs and the corresponding standard deviation. We report in bold the metric that wins between the raw data representation and the augmented representations. We underline the overall best metric for a dataset.
A
HR t
A Baseline Neural Networks
f
CA 4
HI
MNIST
MLP
0.827 #
0.353 #
0.672 =
0.916 2
0.511 =36
0,681 =
0.978 =3.
T-JEPA
0.866 #2.
0.400
0.728 =3
0.961 #0c
0.468 4
0.517 #8c
0.983 #le
DCNv2
0.829
i:e
0.347 dde-
0.662 13c
0.905 13c
0,504 =4c
0,683 133-
0.971 =1c-
+T-JEPA
0.861 =
0.399 =3<7>
0.723 =2
0.955 =2
0.420 #3e-2
0.525 #8c-2
0.981 #2c-3
ResNet
0.814 =
0.351 +2
0.666 +3.
0.919 =
0.534 =2.-
0.674 =
0.979
+T.JEPA
0.865
0.401 =26"
0.718 =3e
0.964
0.441 i:t
0.705
0.983
AutoInt +T-JEPA
0.823
0.338
0.653 iide
0.894 1:2e
0,501 =3e
0.694
0.901 1:0
0.866 #2c
0.351 =3
0.710 =3e->
0.938 =:de
0.448 =:20
0.517 18c
0.978 #1e
FT-Trans
0.821 +7+
0.363 #
0.677 a
0.913 #3 -
0.473 =5c
0.684 +5x
0.811 #5x
+T-JEPA
0.864 =
0.384 =56
0.708 =5e"
0.921 = 0.444 #l
0.551 #00
0.966 ::
Gradient Boosted Decision Trees (GBDT)
XGBoost
0.872 =5
0.375 #1.2c
0.721 =1e-a a
0.951 = -a
0.433 =
0.729 #1e
0,980 =1c
CatBoost
0.873 #1e 0.381 +lc
0.721 #1. 0.946 + 0.430 =7c 0.726 #:e
0.972 #:3
tabular deep learning tasks such as binary and multi-class classification, as well as regression. We use as performance metrics Accuracy (f) and RMSE (L) for classification and regression respectively. The datasets we include in our experiments are Adult (AD) (Kohavi et al., 1996), Higgs (HI) (Van- schoren et al., 2014), Helena (HE) (Guyon et al., 2019), Jannis (JA) (Guyon et al., 2019), ALOI (AL) (Geusebroek et al., 2005) and California housing (CA) (Pace and Barry, 1997). We also add MNIST (interpreted as a tabular data) to our benchmark following Yoon et al. (2020). We summarize the characteristics of all 7 datasets in Table 11 in appendix C.
Baselines To assess whether our self-supervised approach can foster performance 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, 2018), DCNV2 (Wang et al., 2021a), ResNet (He et al., 2016), AutoInt (Song et al., 2019) and FT-Transformer (Gorishniy et al., 2021). Similarly, as often considered as the go-to methods for supervised tasks on tabular data, we also compare the performance of T-JEPA augmented models to XGBoost (Chen and Guestrin, 2016) and CatBoost (Prokhorenkova et al., 2018). To further assess the relevance of T-JEPA to foster performance on supervised tasks, we compare the performance of MI. and ResNet when trained on the original dataspace representations and those generated by PTaRL (Ye et al., 2024), SwitchTab (Wu et al., 2024), BinRecon (Lee et al., 2024), SubTab (Ucar et al., 2021) VIME (Yoon et al., 2020) and T-JEPA.
T-JEPA Training We split each dataset into training/validatiionnest 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 systematic approach to train and evaluate the embedding space generated by T-JEPA. First, following previous work (Assran et al., 2023; Bardes et al., 2022) 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 provided in the supplementary material.
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
6

---

Published as a conference paper at ICLR 2025
Table 2: Comparison of different SSL models with ResNet and MLP as downstream model. See Appendix F for more detail on experimental setting. We report and average over 20 runs and the corresponding standard deviation. The last columns displays the average rank over each dataset, the lower the better.
AD
HE
JA
AL
CA
HT
MNIST
Avg Rank
MLP
0.827 =
0.353 =
0.672 2
0.916 =4-
0.511 =d
0.631 =
0.978 2de
9.4
+PTaRL
0.868 1de
0.383 13e
0.710 14t
0.917
0.489 13e
0.713 13.
0.977 Lit
5.1
SwitchTbb
0.567 aNA
0.387 ANA
0.726 ANA
0.942 ANA
0.452 ANA
0.724 NNA
N/A
4.3
+VIME
0.850
0.362
0.695
0.925
0.505
0.655
0.941 Lit
7.4
BinRecon
0.816 s:s
0,346 #2c
0.581 tie""
0.845 #::
0,498 #ie
0,6016 =he
0,918 =ãe
11.7
+ SubTab
0.823 +
0.378
0,702
0.961 #2:
0.519 =d.
0.673 =2-
0.977 =2u-
8.1
+T-JEPA
0.868 #::
0,400 ::4e
0.728 =3
0.961 2t:
0,468 1:4e
0.517 #
0.983 #ie
3.9
ResNet
0.814
0.351 19.
0.606
0.919
0.534 17.
0.674 1
0.979 =1e
10.4
+PTARL
0.862 15a-3
0.383 -3
0.723 1
0.895 11
0.498 -
0.713 -3
0.973 -
6.1
VTME
0.851
0:372
0.699
0.959
0.505
0.686
0.933
7.6
BinRecon
0.828 1i
0.327 1ite
0.699 1ss
0.944 11.
0.471 1if
0.711 13c
0.981 12t
6.9
SubTab
01822
0.365
=
0,7112 #1e->
0.958 #1.-
0.187 #2c-
0.675 #
0.981
6.3
+T-JEPA
0.865 1.3e
0.401
1.2
0.718 1.s
0.964 Ltx
0.441 1s
0.705 1.c
0.983 12t
2.6
performance, and final evaluations were conducted the 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. We refer the reader to appendix B for details on the downstream models' hyperparameters.
4.2 RESULTS
Downstream Models As displayed in Table 1, T-JEPA pre-training improved the performance of all evaluated models except the Higgs (HI) dataset where only the performance of ResNet improves when augmented by T-JEPA. Interestingly, despite their design advantages for tabular data, recent attention-based models such 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. Interestingly, when compared to gradient boosted decision trees, we observe that augmented with T-JEPA representation, the deep methods regularly obtain the best performance.
Comparison to SSL methods As displayed in Table 2, when compared to existing SSL methods, ResNet+T-JEPA obtains the lowest average rank while MLP+T-JEPA obtains the second lowest average rank. This emphasizes the relevance of T-JEPA for generating representation that foster performance on both regression and classification tasks. Other best performing alternative include ML.+SSiithTaal and MLP+PTTRLL that obtain respectively the third and fourth lowest average rank.
Overall these experiments demonstrate that T-JEPA contributes to increasing the performance of wide range of downstream models. Moreover, appears that T-JEPA outperforms existing SSL methods for tabular data when tested with MLP and ResNet as the downstream model.
5 DISCUSSION
5.1 REPRESENTATION SPACE EVALUATION
Two critical properties are considered desirable in a representation space: uniformitt and alignment (Wang and Isola, 2020). Uniformity measures how well information contained in the original data space preserved. Alignment describes representation space in which semantically close samples are close to one another, while different samples should be distant. An edge case where both uniformity and alignment are not achieved the one of representation collapse (see section 5.2) that describes a situation in which all samples are mapped to the same representation in the latent space.
Metrics We propose two complementary metrics to measure whether our representations are collapsed satisfy the uniformity and alignment properties. First, measure distribution consistency
7

---

Published as a conference paper at ICLR 2025
Table 3: Key metrics tracking T-JEPA's representation learning on the Jannis (JA) dataset over different training epochs.
Metric/Epoch
20
40
60
100
120
DKL 1l112
9.30e -4
3.61e
6.62e -2
5.06e -2
6.55e -2
7.66e
9.38e -2
5.83
3.93
54.3
50.6
68.1
71.4
70.0
un unifor formity
3.12
0.94
10.69
11.20
11.23
11.32
11.38
Epoch 0
Epoch 20
Epoch 40
Epoch 60
Epoch 80
Epoch 100
-1
0
-1
0
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.
and alignment, we rely on the Kullback-Leibler divergence (Dxr.(t)): collapsed representations would imply similar distributions between diverse samples in the embedding space. Thus, lower values would indicate collapse as different samples have indistinct distributions. In contrast, higher values show that dissimilar samples tend to be located farther apart. We expect D increase as training progresses. Second, we rely on the un iniffrrea formioimiiiimmmmmmm tty() metric (Wang and Isola, 2020) to measure how much of the representation space is utilized. Collapsed representations would imply that samples are tightly clustered in a small region of the latent space, achieving low uni f formi ty, On more uniformly spread out, effectively using the embedding
the contrary, desirable representations space. See Appendix G for more detail on the metrics.
Embedding space characterization To assess whether T-JEPA generates representations uniformly spanned across the embedding space, display in Figure the obtained sample representations for the JA dataset at different training epochs. We randomly sample 50,000 points and rely on the PaCMAP (Wang al., 202 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 3 how the uniformity score, KL divergence (DRL), and the euclidean distance (|) (2) vary during training. 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. The increasing KI. divergence and euclidean distance demonstrate that the model effectively avoids representation collapse and better satisfies the alignment property as training progresses. Indeed, T-JEPA's training objective pushes samples farther from one another which facilitates discrimination between samples for supervised tasks. Finally, increasing uni formittyy across epochs is in line with the representations displayed in Figure and demonstrates that the training objective preserves information from the original dataspace.
5.2 REPRESENTATION COLLAPSE
As non-contrastive self-supervised learning method, T-JEPA is prone to representation collapse as discussed in previous work (Bardes et al., 2024; Assran et al., 2023). 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, observed that further regularization of T-JEPA necessary avoid such pitfall. Specifically,
8

---

Published as a conference paper at ICLR 2025
Random initialization
(b) Collapee
150
{(a)
T-JEPA (Ours)
100
(d)
50
Training Loss
(b)
(c) Regularization effects
(d) Convergence
0.3
I-JEPA
0.2
0.1
00
10
15
35
Iteeation:
XIP
Figure 3: Training regime of (ImageNet-IK5 data. We display on the right randomly selected sample's representations for each critical part of the training process. The subfigures (a) to (d) illustrate the evolving outputs of the context encoder Acontext E dd In each heat-map, rows correspond to the d features, while columns represent the h hidden dimensions. (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.
Predictive Architectures tabular (JA) and image
appending a regularization token [REG] to both target and context representations appeared critical to escape the initial representation collapse.
Initial Collapse We observe that JEPA-based models undergo initial representation col- lapse due to the EMA relation between the weights of the context and label encoders. Indeed, retrained from scratch I-JEPA (Assran et al., 2023) on ImageNet- K using their official im- plementation and hyperparameters/ and observed similar training regime as for T-JEPA, First, the loss collapses close 0 after few iterations; then, reg- ularization starts pushing the model's 's weights toward a non-trivial equilibrium. We display in Figure 3 the training regimes of both T-JEPA and I-JEPA.
0 [REG Token
J [REG] Tokon
[RBG] Tobeus
3 [RBG] Toleas
Tran loes
[[4]]]] Tolons
[REG] Tobeus
100
Regularization Token While other JEPA methods do not appear to require further regularization, we ob- serve that appending regularization token [REG] to
0-
Epoch
the sample's representations is instrumental in escap- ing 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, the optimization process can escape this initial collapse and pushes the weights towards a non-trivial equilibrium.
Figure 4: Regularization token ablation. Training loss for the JA dataset across differ- ent numbers of regularization tokens [REG].
5.3 COMPARISON WITH GRADIENT BOOSTED DECISION TREES
Performance Comparison Recent work (Grinsztajn et al., 2022; Gorishniy et al., 2024) discusses how neural networks tend to struggle with structured tabular data type in comparison with other non-deep methods based gradient- gradient-boosted boosted decision trees (GBDT). In most scenarios, approaches such as XGBoost (Chen and Guestrin, 2016) or CatBoost (Prokhorenkova et al., 2018) surpass deep learning algorithms. We observe in Table that on most datasets, methods trained on raw data
For computational purposes, reduced the batch size to 16 and kept unchanged the rest of the hyperpa- rameters.

---

Published as a conference paper at ICLR 2025
significantly outperformed by both GBDT methods. However, augmented by T-JEPA, most methods consistently outperform GBDT or match their performance.
Feature Importance To try and understand why T-JEPA enabled some approaches match the performance of GBDT, we investigated whether high-variance features in the latent space corre- late with feature importance for downstream tasks. While the d representations do not exactly correspond to one-to-one relation with the corresponding feature in the original dataspace, index embeddings allow the obtained encoded representations to still hold feature-related information.
Including index embeddings does not eliminate token mix- modules, nevertheless it
ing from "ssidual/sel--atentiie still helps maintain degree of alignment between features and representations. The embedding variance was com- puted as the standard deviation of each feature's values across hidden dimensions, normalized by the dimension- wise mean. We ranked features by their average embed- ding variance across all samples and compared these rank- ings to those generated by traditional supervised methods, including XGBoost and permutation importance, using Kendall's for rank similarity. Overall, our analysis re- veals strong correlation between high-variance features the T-JEPA embeddings and those identified as impor- tant by supervised methods, despite T-JEPA being trained without target labels. Figure 5 provides a detailed com- parison. The embedding variance ranking (Gembed) shows significant correlation with XGBoost (r = 0.44, with p-value 1.73e-9). 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 aligmment between embedding variance and traditional feature importance further highlights the model's 's ability to learn meaningful data representations.
Random Pubeed Pormitation
0.66
0.44
0.43
0.08
0.03
0.06
XGBoost
Permutation
Gunbbe
Figure 5: Pairwise comparison of feature rankings using Kendall's correlation on the JA dataset. Rankings are derived from XGBoost feature importance, per- mutation importance, T-JEPA's embed- ding variance (Gembea), and random baseline.
CONCLUSION
Overall, we have proposed a novel augmentation--ree 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 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 solid alternative to GBDT for tabular data (Gorishniy et al., 2021: Zabërgja et al., 2024), we observe that ResNet+T-JEPA outperforms all competing methods on most datasets, including GBDT. Finally, empirically uncovered 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 other architectures. Future work may include investi- gating 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 pretext task for self-supervised anomaly detection on tabular data. Finally, our work has emphasized the uncanny training regimes of JEPA-like methods non contraastee self-supervised approaches. Further investigation to provide better theoretical insight into why non-conton non-contrastive self-supervised learning works well might enable the construction of novel designs and approaches that foster performance.
10

---

Published as a conference paper at ICLR 2025
Acknowledgemeent This work performed using HPC resources from the "Mésocentre" comput- ing center of CentraleSupélee and École Normale Supérieure Paris-Saclay supported by CNRS and Région fle-de-Fin Fende-eance (http: /t/////msscet/eeeee central westerreeeseeleee lec. fr/). This work was also granted access the HPC resources of IDRIS under the allocation
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.
REFERENCES
Mahmoud Assran, 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 Computer Vision and Pattern Recognition (CVPR), pages 15619-15629, 6 2023.
Dara Bahri, Heinrich Jiang. Yi Tay, and Donald Metzler. Scarf: Self-supervised contrastive learning using random feature corruption. In International Conference on Learning Representations, 2022.
Adrien Bardes, Jean Ponce, and Yann LeCun. VICReg: for self-super learning. In International Conference on Learning Representations, 2022.
invaaiinnc---cooviiaeee covariance regularization
Adrien Bardes, Quentin Garrido, Jean Ponce, Xinlei Chen, Michael Rabbat, Yann LeCun, Mido Assran, and Nicolas Ballas. Revisiting feature prediction for learning visual representations from video. Transactions on Machine Learning Research, 2024. ISSN 2835-8856.
Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. Advances in neural information processing systems, 33:9912-9924, 2020.
Tianqi Chen and Carlos Guestrin. Xghoost: A scalable tree boosting system. In Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, KDD 16., page 785-794, New York, NY, USA, 2016. Association for Computing Machinery.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In Hal Daumé III and Aarti Singh, editors, Proceed- ings of the 37th International Conference on Machine Learning. volume 119 of Proceedings of Machine Learning Research, pages 1597-1607. PMLR, 7 2020.
Timothée Darcet, Maxime Oquab, Julien Mairal, and Piotr Bojanowski. Vision transformers need registers. In The Twelfth International Conference on Learning Representations, 2024.
Zhengcong Fei, Mingyuan Fan, and Junshi Huang. A-JEPA: Joint-embeddingg predictive architecture can listen, 2024.
Jan-Mark Geusebrock. Gertjan Burghouts, and Amold WM Smeulders. The amsterdam library of object images. International Journal of Computer Vision, 61:103-112, 2005.
Yury Gorishniy, Ivan Rubachev, Valentin Khrulkov, and Artem Babenko. Revisiting deep learning models for tabular data. In Advances in Neural Information Processing Systems, 2021.
Yury Gorishniy, Ivan Rubachev, Nikolay Kartashev, Daniil Shlenskii, Akim Kotelnikov, and Artem Babenko. TabR: Tabular deep learning meets nearest neighbors. In The Twelfth International Conference on Learning Representations, 2024.
Leo Grinsztajn, Edouard Oyallon, and Gael Varoquaux. Why do tree-based models still outperform deep learning on typical tabular data? In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2022.
Isabelle Guyon, Lisheng Sun- Hosoya Mare Boullé, Hugo Jair Escalante, Sergio Escalera, Zhengying Liu. Damir Jajetic, Bisakha Ray. Mehreen Saeed, Michèle Sebag, et al. Analysis of the automl challenge series. Automated Machine Learning, 177:177-219, 2019.
11

---

Published as a conference paper at ICLR 2025
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 6 2020.
Dan Hendrycks, Mantas Mazeika, Saurav Kadavath, and Dawn Song. Using self-supervised learning can improve model robustness and uncertainty. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Noah Hollmann, Samuel Müller, Katharina Eggensperger, and Frank Hutter. TabPEN: transformer that solves small tabular classification problems in a second. In The Eleventh International Conference on Learning Representations, 2023.
Dasol Hwang, Jinyoung Park, Sunyoung Kwon, KyungMin Kim, Jung-Woo Ha, and Hyunwoo J Kim. Self-supervised auxiliary learning with meta- paths for heterogeneous graphs. Advances Neural Information Processing Systems, 33: 01229--13055 2020.
Allan Jabri, Andrew Owens, and Alexei A Efros. Space-time correspondence as a contrastive random walk. Advances in Neural Information Processing Systems, 2020.
Arlind Kadra, Marius Lindauer, Frank Hutter, and Josif Grabocka. Well-tuned simple nets excel on tabular datasets. In Thitt Conference on Neural Information Processing Systems, 2021.
Ron Kohavi al. Scaling up the accuracy of naive-bayes classifiers: A decision-tree hybrid. In Kdd, volume 96, pages 202-207, 1996.
Bruno Korbar, Du Tran, and Lorenzo Torresani. Cooperative learning of audio and video models from self-supervised synchronization. Advances in Neural Information Processing Systems, 31, 2018.
Jannik Kossen, Neil Band, Clare Lyle, Aidan Gomez, Tom Rainforth, and Yarin Gal. Self-attention between datapoints: Going beyond individual input-output pairs in deep learning. In Advances in Neural Information Processing Systems, 2021.
Kyungeun Lee, Ye Seul Sim, Hyeseung Cho, Moonjung Eo, Suhee Yoon, Sanghyu Yoon, and Woohyung Lim. Binning as pretext task: Improving seff-uupervise learning in tabular do- mains. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 26929-26947. PMLR, 21-27 Jul 2024.
Ilya Loshchilov and Frank Hutter. SGDR: stochastic gradient descent with warm restarts. In 5th International Conference on Learning Representations, ICIR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net,
2017 2017.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
Himangi Mittal, Pedro Morgado, Unnat Jain, and Abhinav Gupta. Learning state-aware visual representations from audible interactions. Advances in Neural Information Processing Systems, 35: 23765-23779, 2022.
Daisuke Niizumi, Daiki Takeuchi, Yasunori Ohishi, Noboru Harada, and Kunio Kashino. Byol for audio: Self-supervised learning for general- general-purposs purpose audio representation. In 202/ International Joint Conference on Neural Networks (IJNNN.. IEEE, Jul 2021. doi: 10. 1099/jamn52387.2021 9534474.
Kelley Pace and Ronald Barry. Sparse spatial autoregressions. Statistics Probability Letters, 33 (3):291-297, 1997.
12

---

Published as a conference paper at ICLR 2025
Liudmila Prokhorenkova, Gleb Gusev, Aleksandr Vorobev, Anna Veronika Dorogush, and Andrey Gulin. Catboost: unbiased boosting with categorical features. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
Ravid Shwartz-Ziv and Amitai Armon. Tabular data:
Deep learning is not all you need. In 8th ICML Workshop Automated Machine Learning (AutoML), 2021.
Geri Skenderi, Hang Li, Jiliang Tang, and Marco Cristani. Graph-level representation learning with joint-embedding predictive architectures. arXiv preprint arXiv:2309.1601 2023.
Gowthami Somepalli, Micah Goldblum, Avi Schwarzschild, C. Bayan Bruss, and Tom Goldstein. SAINT: improved neural networks for tabular data via row attention and contrastive pre-trate pre-training.
CoRR, abs/2106.01342. 2021.
Weiping Song, Chence Shi, Zhiping Xiao, Zhijian Duan, Yewen Xu, Ming Zhang, and Jian Tang. Autoint: Automatic feature interaction learning via self-attentive neural networks. Proceedings of the 28th ACM international conference information and knowledge management, pages 1161-1170, 2019.
Yi Sui, Tongzi Wu, Jesse C. Cresswell, Ga Wu, George Stein, Xiao Shi Huang, Xiaochen Zhang, and Maksims Volkovs. self-supervised representation learning from random data projectors. In The Twelfth International Conference Learning Representations, 2024.
Hind Taud and bean-François Mas. Multilayer perceptron (mlp). Geomatic approaches for modeling
land change scenarios, pages 451-455, 2018.
Hugo Thimonier, Fabrice Popineau, Arpad Rimmel, and Bich-Liên Doan. Beyond individual input for deep anomaly detection tabular data. In Proceedings of the 41st International Conference Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 48097-48123, PMLR, 21-27 Jul 2024.
Yonglong Tian, Dilip Krishnan, and Phillip Isola. 05s//996600544, 2019.
Contrastive multiview coding.
CoRR,
Talip Ucar, Ehsan Hajiramezanali, and Lindsay Edwards. Subtab: Subsetting features tabular data for self-supervised representation learning. In Advances in Neural Information Processing Systems, 2021.
Joaquin Vanschoren, Jan N Van Rijn, Bernd Bischl, and Luis Torgo Openml: networked science in machine learning. ACM SIGKDD Explorations Newsletter, 15(2):49-60, 2014.
Ruoxi Wang. Rakesh Shivanna, Derek Cheng, Sagar Jain, Dong Lin, Lichan Hong, and Ed Chi. Den v2: Improved deep & cross network and practical lessons for web-scale learning to rank systems. In Proceedings of the web conference 2021, pages 1785-17977 2021a.
Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through align- ment and uniformity the hypersphere. In Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 9929-9939. PMLR, 13-18 Jul 2020.
Yingfan Wang, Haiyang Huang, Cynthia Rudin, and Yaron Shaposhnik. Understanding how dimen- sion reduction tools work: An empirical approach to deciphering t-sne, umap, trimap, and pacmap for data visualization. Journal of Machine Learning Research, 22(201):1-73, 2021b.
Zifeng Wang and Jimeng Sun. Transtab: Learning transferable tabular transformers across tables. In Advances in Neural Information Processing Systems, volume 35, pages 2902-2915 Curran Associates, Inc., 2022.
Jing Wu, Suiyao Chen, Qi Zhao, Renat Sergazinov, Chen Li, Shengjie Liu, Chongchao Zhao, Tianpei Xie, Hanqing Guo, Cheng Ji, Daniel Cociorva, and Hakan Brunzell. Switchtab: Switched autoencoders are effective tabular learners. Proceedings of the AAAI Conference on Artificial Intelligence, 38(14): 15944)15524415993
Mar. 2024. doi: 10. a0.1660/aaai.v3888 v38114.2953..
13

---

Published as a conference paper at ICLR 2025
Yazheng Yang. Yuqi Wang, Guang Liu, Ledell Wu, and Qi Liu. Unitabe: A universal pretraining protocol for tabular foundation model in data science. The Tweifth International Conference on Learning Representations, 2024.
Hangting Ye, Wei Fan, Xiaozhuang Song, Shun Zheng, He Zhao, Dan dan Guo, and Yi Chang. PTaRL: serotype-based tabular representation learning via space calibration. In The Twelfth Interaational Conference on Learning Representations, 2024.
Jinsung Yoon, Yao Zhang, James Jordon, and Mihaela van der Schaar. Vime: Extending the success of self- and semi-supervised learning to tabular domain. In Advances in Neural Information Processing Systems, volume 33, pages 11033-11043, Curran Associates, Inc., 2020.
Yuning You, Tianlong Chen, Yongduo Sui, Ting Chen, Zhangyang Wang, and Yang Shen. Graph contrastive learning with augmentations. Advances in neural information processing systems, 33: 5812-5823, 2020.
Guri Zabërgja, Arlind Kadra, and Josif Grabocka. Tabular data: Is attention all you need?, 2024.
Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. Barlow twins: Self-supervised learning via redundancy reduction. In International conference on machine learning, pages 12310-12320. PMLR, 2021.
Zehua Zhang and David Crandall. Hierarchically decoupled spatial-temporal contrast for self- supervised video representation learning. In Proceedings of the IEEF/CVF Winter Conference on Applications of Computer Vision, pages 2235-3245, 2022.
Bingzhao Zhu, Xingjian Shi, Nick Erickson, Mu Li, George Karypis, and Mahsa Shoaran. XTab: cross-table pretraining for tabular transformers. In Proceedings of the 40th International Conference on Machine Learning, ICML'223, 2023.
14

---

Published as a conference paper at ICLR 2025
Each experiments detailed in the present work can be reproduced using the following code:
O:
A
T-JEPA TRAINING SETTINGS
This section presents the implementation details the project.
A.1 PROGRAMMING ENVIRONMENT
The code environment for this project was implemented using Python and several third-party libraries. Table 4 details the main libraries used along with their respective versions. The training was done on single NVIDIA HGX A100 GPU with 40GB of memory.
Table 4: Main libraries used in the project.
Library
Description
Python v2.12.2
The programming language used project
einops v0.8, 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_lears v1. v1.4.1.post1
Machine learning library for Python
scipy v1. v1.14.1
Library for scientific and technical computing
torch v2.3.0. postt.30
PyTorch deep learning library
torchinfo v1.8.0
Module to show model summaries 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, 2019) due to its proven robustness in large-scale models. The learning rate was adaptively adjusted using a cosine annealing scheduler (Loshchilov and Hutter, 2017), which gradually reduced it from the initial value to minimum, Mmin O.
Table 5: Hyperparameter Configuration for Bayesian Optimization
Parameter
Values
model_num heads
2, 4,8
2,44, 16, 64, 128 1, 4, 6, 7,8, 16
model__dim_feeeforwwe
[64, 128, 256, 512, 768, 1024]
model_ddzopout_prob exp_1r
(0.0, 0.01)
(0.00001, 0.001)
mask_min_ot__share
(0.07, 0.15)
mask_max_ctx_share
(..2.
mask_min_trgt__haree
(0.05, 0.20)
mask_max_trrt_shhare
(0.2, 0.9)
pred_num_layers
[2, 4. 8, 16, 24, 32| [4, 8, 16, 32, 64, 128 [2, 4.8 (0.0, 0.01)
pred gred_embed_diie dim
pred_num_heade
pred_p_dropoout
The hyperparameters displayed in table correspond to the following:
15

---

Published as a conference paper at ICLR 2025
model_nan num heads: number of attention heads of the context encoder fo-
model_dime hidden: hidden dimension of the context encoder fo- model_num_laye number of layers of the context encoder fa-
model_ding [t____ddd__eeeoorwardod dimension of FFN in the context encoder fa.
model model_drroout__prok dropout probability of the context encoder So-
exp_ir: learning rate.
Minimum share masked feature for the reppesstipttnnntaaiionn (see (a) in Figure 1).
mask_max_ _____ax___tt__ Boii__tt__s_______ share:
_aax_tt_sshare: Maximum share of masked feature for the context representa- tion (see (a) in Figure 1).
mask mask_min__ saii__trg_ssharee Minimum share of masked feature for the target representation (see (b) Figure 1).
mask_max_taae pe_.ttrggt_hhreeee Maximum share of masked feature for the target representa- tion (see (b) in Figure I).
pred_num_headee number of attention heads of the predictor gø.
n_m mw__nmum__yesrs number of layers of the predictor 9é-
pred_embed_dimm hidden dimension of the predictor 9o-
pred p_deppodr dropout probability of the predictor go-
A.3 PROJECTION LAYER
The projection layer adapts the T-JEPA representation space h E Rddh to the input dimensions required by the downstream models. Several projection techniques were implemented, as described below:
When using the linear flatten projection, the input is flattened into single vector hgaaten E Ax and transformed through linear projection, hpoj = W- W-hnaeen + RR/ x1, with W E Rhe xd-h the weight matrix and E R x1 the bias
Linear Flatten:
Linear Per-Feature: Each feature transformed independently by applying a linear projection each feature vector h, gdx1 hr W, W,, R, where W; Rixd is the weight matrix and b, E IR 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 x E Rdx where d the number of feature and h is the hidden dimension, the convolution operation is represented as: x' a (BatchNorm (Conv2D(X. ki))), where ki convolutional kernel, and a is the activation function (ReLU). After second convolution and pooling step, the final representation is flattened into vector and projected to the target embedding dimension.
Max Pooling: The max pooling operation selects the maximum value for each feature vector h, E Rh. h max(h;,) [d))
Mean Pooling: The mean pooling operation averages values for each feature vector h, E Rh: hi.mean Σ h.j-
Each projection layer is trained jointly with the downstream task and is selected based the structure of the input data and model requirements.
B DOWNSTREAM MODELS' ARCHITECTURE
Similarly T-JEPA, employed Bayesian optimization to tune the hyperparameters of the down- stream model's architecture. We give details on each architecture and their corresponding hyperpa- rameters in the present section.
16

---

Published as a conference paper at ICLR 2025
B.1 MLP
The Multi-Layer Perceptron (MLP) architecture employed in this work is designed to effectively transform input features. The transformation begins with linear projection of the input z, which can either be the raw input x or an embedding projection, into a hidden representation dimensionality h:
h (0)
Wiz + b
The model then applies the hidden layers sequentially:
h (++1) BatchNNrm (Dropout [ReLU(Wr.: h h + bi+1))),
where 0, ..... L Here is the selected nyper-parameter for each dataset:
Table Hyperparameters of MLP Model for Each Dataset
Dataset
Dropout
Encoder Type
Learning Rate
Weight Decay
Hidden Layers
HE 0.5478
Semear_ffatter
8.97
10 2
4.45 10-4
4
HI
0.4203
Linear_per_featurr
1.84 10-5
7.48
10
9
JA
0.4923
1inear_per_featur
7.31 10
2.09 10-6
13
AD
0,2311
1.35 10
6.43 10
13
0.0310
linear_flater flatten
1.19 10-5
1.96
10-5
3
AL 0.2331
linear_ limear_flatten flatten
1.87 x 10-4
1.38 x 10
4
MNIST
0.3527
linear_flatten
1.83 10-5
1.47 10
5
B.2 IMPROVED DEEP CROSS NETWORK
This section presents the enhanced DCN-V2 model architecture, designed to learn both explicit and implicit feature interactions. The code used in this work was taken from (Wang et al., 2021a). The DCN-V2 model combines Cross Network with Deep Network, achieving superior expressiveness while maintaining computational efficiency. The explicit feature interactions are modeled through the Cross Network layers, defined as:
X|+1 X C (Wix: + b)) +x, Wi E Rd b E Rd.
Here, Xy represents the base features, and @ denotes element- wise multiplication. Implicit feature interactions are captured by Deep Network, where the I-th layer is defined as:
h f(W, (WWh++bb) b:), f(-) ReLU(-).
Table 7 summarizes the hyperparameters used in experiments.
Table Hyperparameters for DCN-V2 experiments.
Dataset
Cross Dropout
Embedding Dim.
Hidden Dim.
Hidden Dropout
Learning Rate
Weight Decay
HE
0.0808
128
768
0.2926
8.81 x 10-5
9.33 x 10-4
III
0.0346
488
0.0879
5.53 x 10
3.64
JA
0.0702
386
0.0862
8.79 x 10-5
1.21 10
AD 0.1393
428
0.2976
4.60 x 10-6
1.21 10-4
CA 0.2578
0.1111
10
10
AL 0.1476
524
0.0431
1.94
10-5
1.40 10-6
MNIST
0.2836
755
0.0875
4.22
10
1.77
10-5
B.3
RESNET
The ResNet model used in this work was specially tailored for tabular data. The model is formalized follows:
ResNet(x: Prediction(eeSNeBBlockk.... (esbeeoee
ResNetBlock(ec) = x + Baapauuinnneaannonntneteeeee (BatrrrrnnnnNNo()))))
(2)
Prediction(2) LLiner(Re.U((BacNhNommm)))
17

---

Published as a conference paper at ICLR 2025
Table Hyperparameters for ResNet Variants
Dataset
block
dropout1
dropout2
Ir
blocks
HE
512
0.459
0.461
2.74 x 10
AD
440
0.042
0.137
8.90 x 10
CA
465
0.313
0.002
5.26 x 10-5
HI
506
0.234
0.031
2.24 x 10-5
AL
476
0.044
0.049
1.32 10 5
JA
355
0.137
0.005
4.60 X 10
MNIST
497
0.265
0.050
6.25 X 10
B.4 AUTOINT
AutoInt a neural network leveraging self-attention mechanisms for automatic feature interaction learning. The input features E R" are embedded into vectors E that interact through multi-head attention. These interactions produce ierraction-adjuster embeddings e, A residual connection ensures the retention of low-order information:
Res ReLU(ē; + WRese))
The final prediction computed as:
o ( (eRcs + ... @ N b),
where a (I) represents the sigmoid function.
The hyperparameter configurations for different datasets are summarized in Table 9,
Table 9: Hyperparameter configurations for the AutoInt model across datasets.
Dataset
token
Num. layers
Learning rate
Residual Dropout
Weight Decay
AD
200
1
1.717 x 10-3
0.060
1.782 x 10
CA
238
8
2.141 x 10
0.071
1.892 x 10
HI
166
1
6.913 10-4
0.043
1.385 10-7
AL
228
3
2.280 x 10-3
0.011
4.667 x 10
JA
32
6
9.680 x 10-4
0.090
1.234 x 10
HE
128
6
5.193
10-5
0.055
7.590 10-4
MNIST
252
5.831 2
X 10-4
0.091
5.107 X 10-5
B.5 FT-TRANSFORMER
The FT-Transformer processes tabular data using three key stages: feature tokenization, sequential Transformer layers, and prediction head. Feature tokenization encodes each numerical feature x; and categorical feature E, as:
TP(uu)) (mum)
xr Wr(mm))
Trlcat)
1/(am)) + e, We(er)
These embeddings are concatenated into a matrix E Rxd where le is the number of features and is the embedding dimension.A [CLS] token is prepended to T, and L Transformer layers are applied, iteratively updating the representation as follows:
T MHSA(LN(T)-))) + T.-1. T, FFN(LN(T))) + T,
where LN is layer normalization. The prediction head computes the final output by processing the [CLS] token with normalized and activated linear transformation:
i W₂ ReLU(W, LayerNorm(71..1)))
The following table summarizes the hyperparameters for the FT FT-Trane Transsonnfrrrss model applied to various datasets.
18

---

Published as a conference paper at ICLR 2025
Table 10: Hyperparameter Configuration per Dataset
Dataset
Att. Dropout
Num. Heads
Block Size
Hidden Dim
Learning Rate
HE
0.0157
16
64
128
1.08 x 10-3
HI
0.3334
4
128
64
2.25 10
JA
0.0934
16
64
256
1.33 x 10-1
AD
0.3426
16
128
32
1.35 x 10-1
CA 0.0745
8
128
512
1.17 10-3
AL 0.3246
8
256
64
8.27 X 10
MNIST
0.4259
8
32
256
7.97 X 10-5
C COMPUTE AND TRAINING TIME
The training process executed 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 11, 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 balance between model performance and resource usage.
Table 11: Dataset characteristics and pretraining GPU-hours.
AD
HI
JA
AL
CA
MNIST
Samples
48,842
98,050 28, 0 2
65,196
83,733
108,000
20,640
67112
Numerical, Categorical
6,
27, 0
54, 0 4
128, 0
8, 0
784, 0
Classes
2
100
1,000 Accuracy
N/A
10
Metric
Accuracy
Accuracy
Accuracy
Accuracy
RMSE
Accuracy
Pretraining GPU-hours
0.84
1.80
1.03
1.27
3.64
0.34
4.80
D 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 D. Second, we also considered an alternative where the masking strategy is the one detailed in section D.1, and we also modify the predictor architecture as detailed in section D.2.
D.1 MASKING STRATEGY
Masking Following previous work on feature masking for tabular data, we considered handling masked features by keeping the sample's 's representation's dimension constant. Features of a samples normalized similarly detailed in section such that E(x,) RS, where e, for numerical features and for categorical features éj corresponds to their cardinality. Each sample is accompanied by masking vector m E {0, 1] in which each entry designates whether a feature is masked: m, 1 ffetuure j 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 (ej di-dieensio representation
E(x) = ((1g - m) @ E(x), m) E Idx(ee+1),
(6)
where @ designates the Hadamard product and 1 the d-dimensional unit vector.
We pass each of the d features encoded representations of sample through learned linear layers Linear(e, + 1, h). We also learn h-dimensional index and feature-type embeddings. Both are added to the embedded representation of sample x. Let z E d denote the obtained embedded representation of sample with mask
19

---

Published as a conference paper at ICLR 2025
Context and Target Encoders Given this modification, the obtained context representation's dimension differ from the one given in equation (2),
1 context
fa E Rdd (context)
The rest of the pipeline is identical to the one given in section We also considered second alternative where the above pipeline is chosen but different predictor architecture (see section D.2) replaces the transformer as detailed in section
D.2 PREDICTOR
Given this alternative masking strategy, we also considered MLP-based predictor for target representation prediction. The predictor consists of separate MLPs (one for each feature). Each MLP takes as input representation of dimension d' (the flattened representation output by the context encoder), and produces and output of dimension h. In particular, each MLP corresponds to one feature particular and produces an I--dimensional prediction for the corresponding feature given a flattened context representation of dimension d h.
Let us denote O/ = [MLP:] a target mask migt and 2flatten tentten(( Then the prediction for the target mask marget is given by,
Mecrggt (MLP) (EFLat(E))atten) m = 0}
E EMBEDDING FEATURE VARIANCE
The idea behind calculating the variance of embedding features based on the assumption that features with high variability across the embedding space 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).
High-variance features
Gmbee
Realure trres
211
Hidden features
Figure 6: Embedding Variance
To measure this variance, let us define the embedding variance score. Let x E dd represent 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 μj across the d features, i.c., μy = - 1 Ti.j. For each feature we calculate the embedding variance score, defined as d'.embedd Var(x; 14), where x; is the vector representing the i-th-feature and [Mo.... 1", "", Features with higher embedding variance are those that stand out more in the representation space, suggesting they may carry more relevant information.
20

---

Published as a conference paper at ICLR 2025
F EXPERIMENTAL SETTING
In Table I, we present performance metrics for several models, including neural networks (MLP, DCNv2, ResNet, AutoInt, FT-Trans), and with and without T-JEPA. For both cases, we relied on bayesian optimization to find the best set of typerparameters according to performance on validation set. We report the performance obtained with this set of hyperparameters on test set never used during training.
In Table we report the performance of several SSL methods (PTaRL (Ye al., 2024), SwitchTab (Wu et al., 2024). BinRecon (Lee et al., 2024) and SubTab (Ucar et al., 2021)). Metrics displayed in this table are obtained as follows:
For PTaRL (Ye al., 2024), report the metrics from their paper for datasets AD, JA and CA as the authors kindly shared the standard deviations corresponding to these datasets for this model. For the remaining datasets (HE, AL, HI and MNIST) display metrics obtained from our own experiments and their corresponding standard deviations.
For BinRecon (Lee et al., 2024), metrics for MLP are primarly obtained from their paper, while the experiments ourselves for the ResNet alternative.
For Tab (Wu et al., 2024), metries are obtained from their paper for the MLP dosnwtream model, except for MNIST as they did not use this dataset in their benchmark. Moreover, as the authors have not released any code for their method, are unable to obtain any metric for the MNIST dataset and the ResNet downstream model. Moreover, no standard deviations reported in their paper.
For VIME (Yoon et al., 2020), we used the official code made available online by the authors to run the experiments for HE, JA, AL, CA and HI datasets and report the metrics from their paper for AD and MNIST for the MLP downstream model, and the experiments for all datasets for the ResNet downstream model. Notably, we rely on the self-supervised only set-up for fair comparison with other methods.
For SubTab (Ucar al., 2021), we rely on their official implementation and run experiments on each of the datasets except for AD which we obtained from their paper for the MLP downstream model.
For each SSL model, downstream model and dataset we rely on bayesian optimization to set the hyperparameters. This ensure fair and reprodueible comparison between all methods. We provide the code to replicate all our experiments on the official github repository and we display in appendix F. and F.2 the experimental details.
BASELINE PERFORMANCE
The default hyperparameters for each model were utilized ensure consistent configurations and establish baseline performance metrics. To maintain alignment with the experimental framework presented in Ye et al. (2024), the same hyperparameters were adopted for generating the baseline results.
21

---

Published as a conference paper at ICLR 2025
Table 13: DCNV2
Default Hyperparameters for
Table 12: Default Hyperparameters for MLP
Hyperparameter
Value
Hyperparameter
Value
Hidden Dimension
128
Number of Layers
4
Number of Cross Layers
Hidden Dimension
256
Number of Hidden Layers
7
Dropout
0.1
Cross Dropout
0.1
Batch Size
128
Hidden Dropout
0.1
Learning Rate
1 x 10
Batch Size
128
Early Stopping Patience
16
Learning Rate
10-4
Maximum Epochs
200
Early Stopping Patience
16
Categorical Embedding Dim
128
Maximum Epochs
200
Categorical Embedding Dim
128
Table 15: Default
Hyperparameters for
Table 14:
Default Hyperparameters for
AutoInt
ResNet
Hyperparameter
Value
Hyperparameter
Value
Hidden Dimension
192
Number Layers
Number of Layers
Hidden Dimension
256
Number of Heads
8
Hidden Dropout
0.1
Attention Dropout
0.1
Batch Size
128
Residual Dropout
0.1
Learning Rate
1 x 10-1
Batch Size
128
Early Stopping Patience
16
Learning Rate
1 10-4
Maximum Epochs
200
Early Stopping Patience
16
Categorical Embedding Dim
128
Maximum Epochs
200
Table 16: Default Hyperparameters for FT-Transformes
Hyperparameter
Value
Hidden Dimension
192
Number of Layers
Number Heads
Attention Dropout
0.1
Residual Dropout
0.0
Batch Size
128
Learning Rate
x 10-1
Early Stopping Patience
Maximum Epochs
200
F.2 SELF-SUPERVISED METHODS
SwitchTab SwitchTab is framework for tabular data representation learning that utilizes anencoder-decoder architecture to decouple features into mutual (shared across samples) and salient (unique individual samples) representations. The process begins with feature corruption applied to input data to enhance robustness. The corrupted data is then encoded by an encoder, producing feature vectors that are decoupled into mutual and salient components using projectors. These components are recombined and reconstructed by a decoder, with both recovered and switched outputs contributing the computation of reconstruction loss.
The results presented are drawn from the original work (Wu et al., 2024). Since the code was not available, results for the MNIST dataset were not included.
22

---

Published as a conference paper at ICLR 2025
BinRecon In the approach presented in (Lee al., 2024), continuous numerical features discretized into a fixed number bins, where each bin represents a range of values defined by quantiles of the training dataset. The task is predict the bin indices instead of reconstructing the raw values, effectively transforming the problem into regression or classification task depending whether the bins treated ordinal categorical. This encourages the encoder learn representations that capture irregularities and nonlinear dependencies characteristic of tabular data. The loss for BinRecon, when treating bins as ordinal values, is defined as:
BinnReco Recen
144
Bin /einneeem())]13
where i; denotes the bin indices of the input features, 2, represents the encoder outputs, and Bin is the decoder network.
VIME
The VIME framework a self-supervised learning method for tabular data that leverages two pretext tasks: feature estimation, which involves reconstructing the original features, and mask estimation, which focuses on predicting the applied binary mask. masked sample x generated as:
m + x + (1 m) O x,
where m is binary mask, and x values sampled from marginal distributions. An encoder maps X to which is used to predict the mask (m) and reconstruct the input (x). The framework minimizes:
min E Im (m, m) + a (x, x)],
where im is binary cross-entropy and 1, is reconstruction loss.
The datasets AD and MNIST were sourced from the work presented in (Yoon et al., 2020). For the remaining datasets, the publicly available code repositories were utilized, ensuring reproducibility and adherence to standardized evaluation protocols.
The hyperparameters presented Table were obtained following hyperparameter tuning pro- cess, designed to optimize performance metrics for each specific dataset. This tuning involved Bayesian hyper-parameter search across a range values for parameters such as alpha, beta, halp_hidden, _imm and p_m, ensuring that the reported results reflect the best possible configura- tions for the proposed approach.
Table 17: Hyperparameter Settings for VIME experiments
Dataset
alpha
beta
dim
p_m
JA
3.1343
1.5921 x 109
32
1.5536 10-1
AL
1.3867
1.0233 x 10"
256
2.4229 x 10
HE
4.7568
8.1190
10-1
256
1.2751 10-1
HI 4.7680
1.4418 x 109
128
2.4642 10-1
CA
2.5088
256
1.1916 x 10-!
SubTab SubTab a framework for representation learning on tabular data inspired by cropping in image augmentation. It divides tabular data into subsets of features, processed by a shared encoder-decoder architecture.
The results in Table were derived using the AD and MNIST datasets, sourced from (Ucar al., 2021). For the other datasets, publicly available code repositories were employed, ensuring reproducibility.
The hyperparameters listed in Table 18 were determined through hyperparameter tuning pro- cess aimed at optimizing performance metrics for each dataset. This process utilized Bayesian hyperparameter search across a predefined range of parameter values.
23

---

Published as a conference paper at ICLR 2025
Table 18: Hyperparameter Settings for SubTab experiments
Dataset
Dropout Rate
Hidden Layers
Learning Rate
Masking Ratio
N Subsets
HE
0.1941
[1024, 256]
1.0867
10
0.2
4
JA
0.1383
[1244, 1024, 128)
1.6595 10 3
0.3
4
0.1542
[1024, 512]
6.1929
10
0.1
CA 0.0292
[128, 128]
5.8552
10-4
0.1
HI 0.1310
[122, 512]
1.4141
10-3
0.2
4
G REPRESENTATION SPACE CHARACTERIZATIIN
Kullback-I Leibler divergence The Kullback-Leibe (KI.) divergence measures the difference between the probability distributions random points within the embedding space. Formally, given two probability distributions P and Q over the same variable æ, the KL-divergence from Q to P is defined as:
P(r) Q())
(p = Σ P(x) log
(8)
We consider random points z, =fatten(h;) and 2, = flatten(hh) from the embedding space and estimate their distributions P and Q, where flatten is operation that converts the high- dimensional embeddings E Rd into a one-dimensional vector x E IRd-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, 2020) to evaluate the preservation of maximal information within the feature distribution. This score leverages the Gaussian potential 0. The uniformity score is defined as the logarithm of the average pairwise Gaussian potential, formally:
kernel G, Sd S + R,,, defined (u, v) & e e-1
uniformity & log
[G) log
E
ti t> O.
(9)
a.
prn
This metric is intricately connected to the notion of uniform distribution the unit hypersphere. This uniformity score allows us to obtain a nuanced measure that captures the degree of information preservation in the feature distribution.
24

---

Published as a conference paper at ICLR 2025
H
REGULARIZATION TOKEN
Context encoder
Context
Predictor
Targets pred.
(a)
(d)
fo
96
Original sample
[ECC]
(b)
(c)
(c)
fa
L2
Masked target
[EG)
Target encoder
Target
Unmasked target
Figure 7: Regularization Token Pipeline. In step (a) sample Rd is pre-processed and masked according to context mask, a [REG] token is appended to the masked representation. In step (b) the whole unmasked sample is pre- pre-processed and fed to the target encoder, including the [REG] token. In step (c) the output of the target encoder is masked according to a target mask (see equation 4). fed the predictor to predict the target output. In step (c) we compute the fa-distance between the masked target representation and the corresponding prediction.
In step (d). the [REG] token discarded and the remaining features' representations
As discussed in sections and 5, including regularization tokens in the T-JEPA training pipeline is instrumental to escaping representation collapse. While EMA between the weights of the context and target encoders is sufficient for JEPA-based approaches for images (Assran et al., 2023) or videos (Bardes al., 2024), we show in section 5.2 that without [REG] tokens, the training loss stuck in a collapsed equilibrium.
We provide this section a more detailed description on how [REG] are involved in the training pipeline. In particular, we provide in Figure descriptive example on how the regularization tokens handled.
[REG] Token Pipeline Recall that 2 E Rluxh denotes the obtained embedded representation of sample x with mask m.
Context: Let m. denote the context mask. The context representation is obtained by concatenating the masked embedded representation and the embedded [REG] token, zm concat(z".. [REG]), and passing it through the context encoder. One obtains
A Am Aeontext
fo (2 E ((m,, +1)xh
where lm. corresponds to the number of unmasked features mask m., and the supple- mentary dimension comes from the [REG] token.
Target: The target representation is obtained by concatenating the unmasked embedded representation and the embedded [REG] token, z0. concat (z [REG]) and passing it through the target encoder. One obtains
hrarget = fa(2 E (d+1)xx
Then this representation is masked detailed in equation using the target mask m, to obtain Amger- Note that the [REG] token is always masked/dropped here. See step (c) in Figure 7.
25

---

Published as a conference paper at ICLR 2025
The [REG] token is dropped from the context representation that serves to predict the target representation. One thus obtains context-
m. context
drop
(RSG)
Nooneext E. R/ Xh
See step (d) in Figure
Prediction: One then uses i Mo
predict Iugget
target
go conteesss mc).
Note that, while we experiment using more than one [REG] token in section 5.2, the rest of the experiments presented in this work include one [REG] token.
26


Reviews:
[{"review_id": "aroHm4pkA4", "reviewer": "Reviewer_xkk6", "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. \n\n[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.", "questions": "1. What is $h$-dimensional index and feature-type embeddings, defined on Line 195? \n2. 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? \n3. For Equation(5), I think there is a missing superscript $j$ as $\\mathbf{m_j} \\in M_{context}$. \n4. 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.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "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.", "weaknesses": "I am not convinced about the contributions of Sections 5.1 and 5.3. These sections look artificial and drawn-out to me. \n\n## Comments on Section 5.1\nThe 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\n> Line 366 \"A lower KL-divergence indicates that the embedded space has consistent and similar distributions for different regions.\"\n\n> Line 377 \"This uniformity score allows us to obtain a nuanced measure that captures the degree of information preservation in the feature distribution.\"\n\n> Line 422 \" stable intra-feature variance that indicates consistent information extraction, while the decreasing inter-feature variance shows improved feature specialization\"\n\nare imprecise and hand-wavy. If anything, these arguments confuse me more than provide any explanation. \n\n## Comments on Section 5.3\nI 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TCao0n4u42", "reviewer": "Reviewer_eHM5", "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.", "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].\n2. 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.\n3. 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?\n4. I suggest the authors include the results comparison with PTaRL to table 1, as the two methods are all enhancements to table learning paradigm.\n\n[1] Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture. CVPR 2023.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "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.\n2. 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.\n3. 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.", "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.\n2. 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.\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bmEpb9Zd3V", "reviewer": "Reviewer_NZ9j", "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.", "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.  \n2. 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.  \n3. 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.  \n4. 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.  \n5. The description about right four subfigures of Fig.3 is not enough.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed SSL method is model-agnostic and could be coupled with various deep tabular models, which has wide range application.  \n2. The strategy for masking representations is novel in tabular domain.  \n3. The experiment is extensive.", "weaknesses": "Please refer to the questions below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0LaVsFjZJm", "reviewer": "Reviewer_haUm", "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.", "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.\n\n- 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.\n\n- What are the variances in the reported results?\n\n- It is essential to include comparisons with other SSL methods for tabular data in the main paper.\n\n- Given recent advancements in LLM-based methods and foundation models, it might also be relevant to consider these approaches for this problem.\n\n\n-----------------------------------------------------------------------\n**After Rebuttal:** \nThank 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.\n\nUnfortunately, I am unable to provide further comments on the paper. However, if the authors are able to respond, it would be greatly appreciated:\n\n1- 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?\n\n2- 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.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The method is simple and straightforward, addressing an important question.", "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.\n\n- 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?\n\n- 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.\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)