Summary
The authors propose a self-supervised representation learning approach for tabular data based on the Joint Embedding Prediction Augmentation (JEPA) framework. Methodologically, the model is based on the work of Assran et al. (2023): inputs are masked, passed to 2 encoders, and a prediction model is trying to reconstruct one of the encoded inputs from the other. Empirically, the authors show that using the features from their model, T-JEPA, increases performance on the downstream tasks. The authors also offer some discussion, highlighting that the identified features are indeed meaningful.
[1] Assran, Mahmoud, Quentin Duval, Ishan Misra, Piotr Bojanowski, Pascal Vincent, Michael Rabbat, Yann LeCun, and Nicolas Ballas. "Self-supervised learning from images with a joint-embedding predictive architecture." In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 15619-15629. 2023.
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.
## Comments on Section 5.1
The proposed metrics in Section 5.1 have little to do with the optimized objective, Equation 5. Therefore, I have no intuition of what we want these metrics to be: what is their range? what values are good / bad? how fast / slow do we want them to converge? why do we think optimizing our objective promotes this? The authors' arguments such as
> Line 366 "A lower KL-divergence indicates that the embedded space has consistent and similar distributions for different regions."
> Line 377 "This uniformity score allows us to obtain a nuanced measure that captures the degree of information preservation in the feature distribution."
> Line 422 " stable intra-feature variance that indicates consistent information extraction, while the decreasing inter-feature variance shows improved feature specialization"
are imprecise and hand-wavy. If anything, these arguments confuse me more than provide any explanation.
## Comments on Section 5.3
I am not sure I fully grasp what the author do in Section 5.3. As far I as I understand, they want to compare the feature importance produced by XGBoost (and permutation test) to the feature importance produced by their model. However, I don't think the quantity that the authors claim to be feature importance for their method, which is the variance of the hidden features, is a meaningful proxy. The main problem with this metric is the following. Let's say we have feature vector, $\[x_1, x_2\] \in \mathbb{R}^2$, which we encode to get $\[\mathbf{h}_1^T, \mathbf{h}_2^T\]\in\mathbb{R^{2\times h}}$. The authors propose to compute variance of $\mathbf{h}_i$ across hidden dimension and use this quantity to rank $x_i$. However, in general, $\mathbf{h}_i$ does not correspond to $x_i$. Specifically, in the transformer model, when we compute the self-attention, each of $\mathbf{h}_i$ becomes a weighted some of all the other tokens. In the very simple case, consider what happens, with hidden dimension $h=1$ and if the learned encoding is a simple permutation. Clearly, this is not correct.
Questions
1. What is $h$-dimensional index and feature-type embeddings, defined on Line 195?
2. I think Figure 1 does not quite match the description in Section 3. In particular, in Figure 1, I get a sense that there is a single context mask for each sample, while on Line 216 and Line 241, there appears to be multiple context mask for a single sample. In general, how are these masks drawn and how do you ensure that the context / target masks don't overlap?
3. For Equation(5), I think there is a missing superscript $j$ as $\mathbf{m_j} \in M_{context}$.
4. I don't quite follow why we can't do masking for target features before encoding, the same as we do for context. Specifically, this creates data leak, and I suspect that the mode collapse that you demonstrate in Figure 3 is related exactly to this. For a simplified example of why the data leak happens, consider that we have context and target masks, $\mathbf{m}_1$ and $\mathbf{m}_2$ respectively. Further, for simplicity, suppose $\|\mathbf{m}_1\| = |\mathbf{m}_2|$. Then I claim that picking $f\_\theta$ to be the identity map and $f\_{\bar\theta}$ to be the appropriately constructed permutation, we can achieve 0 reconstruction error. For the context features, we have the following transformations: $\mathbf{z}\_{\mathbf{x}} \xrightarrow{\mathbf{m}_1} \mathbf{z}\_{\mathbf{x}}^{\mathbf{m}\_1} \xrightarrow{f\_\theta} \mathbf{z}\_{\mathbf{x}}^{\mathbf{m}\_1} $. For the target, features we have: $\mathbf{z}\_{\mathbf{x}} \xrightarrow{f\_{\bar\theta}} \[\mathbf{z}\_{\mathbf{x}, \pi(1)}, \mathbf{z}\_{\mathbf{x}, \pi(2)}, \dots, \mathbf{z}\_{\mathbf{x}, \pi(d)}\] \xrightarrow{\mathbf{m}_2} \mathbf{z}\_{\mathbf{x}}^{\mathbf{m}\_1}$, where $\pi$ is the constructed permutation. My conjecture is that the models initially learn this, but then are thrown off by the additional regularization. However, overall, I am not sure if the formulation is correct.