[NeurIPS 2024 Series] Attention-Based Data Augmentation Techniques to Enhance Performance of Tabular Data Prediction

At NeurIPS 2024, the Table Representation Learning Workshop (TRL), held for the third time this year, allowed researchers to focus on the representation learning of tabular data to exchange ideas. This year's TRL workshop saw 61 research papers accepted, a 1.5x increase from last year’s 41 papers, highlighting the growing attention being given to table data research. Table data is one of the most commonly used data formats in industry and everyday life, and global companies such as Amazon and Microsoft are actively conducting research in this area.

At the NeurIPS 2024 TRL Workshop, LG AI Research's Data Intelligence (DI) Lab published a paper titled “AGATa: Attention-Guided Augmentation for Tabular Data in Contrastive Learning,” which was noted for its novel approach to improving training performance while maintaining the structural features of tabular data.

LG AI Researchs’ AGATa: Attention-Guided Augmentation for Tabular Data in Contrastive Learning^[8]

Tabular data is a core data format in various industries such as finance, healthcare, and manufacturing. LG AI Research and manufacturing plants are utilizing tabular data to improve quality control, equipment maintenance, and production efficiency. While tree-based algorithms such as XGBoost^[1] have traditionally been the dominant way to process tabular data, deep learning models based on self-supervised learning, especially contrastive learning, have been gaining traction in recent years^[2][3][4].

However, tabular data contains hundreds of features and complex interactions that, if not handled effectively, can lead to errors in critical decision-making. Due to the nature of tabular data, which unlike images or text has no spatial or sequential structure, applying traditional data augmentation techniques runs the risk of distorting important relationships in the data or damaging its meaning. For example, traditional techniques such as randomized Masking^[5], Shuffling^[6], or CutMix^[7] can overlook important interactions between features in tabular data, leading to poor model performance. Given the structured and heterogeneous nature of tabular data, the effective application of contrastive learning requires novel techniques that enable meaningful augmentation while preserving the core structure of the data.

To address these issues, LG AI Research has developed a new approach called Attention-Guided Augmentation (AGATa), which aims to maximize the performance of contrastive learning by performing meaningful augmentation while preserving the core structures fo tabular data This opens up the possibility of achieving high prediction performance even on complex tabular data that has been generated in real-world settings, such as manufacturing plants, further increasing the practical value of our study.

 

Image 1. A Schematic Diagram of AGATa
(a) Schematic Diagram of Self-Supervised Pre-training with AGATa Augmentation Technique Applied; (b) AGATa Augmentation Module^[8]

During training, AGATa performs augmentations before feeding data into the model, dynamically introducing changes at each epoch to maximize data diversity. It leverages self-attention scores from the Transformer model to identify less important features and randomly applies one of three augmentation techniques: Masking, Shuffling, or CutMix. The augmented data is then used as positive pairs for contrastive learning, with different transformations applied at each epoch to ensure the model learns from diverse variations of the data.

1) Attention-Guided Feature Selection

AGATa uses the attention score obtained from the self-attention mechanism of the transformer model to compute the importance of each feature.

For a given batch B, number of heads H, and number of features F, the attention matrix $A\mathit{\in }{ℝ}^{\mathit{B}\mathit{x}\mathit{H}\mathit{x}\mathit{F}\mathit{x}\mathit{F}}$ is defined as follows:

$A=softmax \left(\frac{Q{K}^T}{\sqrt{d_k}}\right)$

Here, Q and K are the query and key matrices, and, $d_{\mathit{k}}$ is the dimension of the key.
We compute the feature importance matrix $A_{mean} \in {ℝ}^{BxF}$ by averaging the attention scores across all heads.

$A_{mean}=\frac{1}{H}\sum _{h=1}^H{A}_h$

Then, we compute the average importance of each feature to obtain a feature-specific importance vector  $A_{feature} \in {ℝ}^F$.

$A_{feature\left[i\right]} = \frac{1}{F}\sum _{j=1}^F{A}_{mean}[:,i,j]$

Low-importance features are defined as features whose attention scores are in the bottom k%, and k is set to 40% in AGATa. This way, we can preserve the core structure of the data while changing less important features to increase the generalization performance of the model.

2) Dynamic Augmentation Strategy

AGATa maximizes the diversity of data by randomly applying three augmentation techniques—masking, shuffling, and CutMix—to less important features. Each technique is defined as follows:

(a) Masking : For the selected feature $S\subseteq \{1,2,...,F\}$, we replace its value with the mean ${\mu }_{\mathit{k}}$ of the dataset or a fixed value.

$x_i^{\mathit{\text{masked}}}\left[k\right] = \{\begin{array}{l}{\mu }_k,\\ x_i\left[k\right],\end{array}\begin{array}{r}\mathit{\text{if }}k\in S,\\ \mathit{\text{otherwise}}.\end{array}\}$

(b) Shuffling: We replace the values of the selected feature with the values of another sample $x_{\mathit{j}}$ in the batch.

$x_i^{shuffled}\left[k\right]=\{\begin{array}{l}{x}_j\left[k\right],\\ x_i\left[k\right],\end{array}\begin{array}{r}if k\in S, and j \ne i\\ otherwise.\end{array}\}$

(c) CutMix: We generate the selected feature values as a mixture of the two samples $x_{\mathit{A}}$and $x_{\mathit{B}}$.

$x_i^{cutmix}\left[k\right]=M\left[k\right]·x_A\left[k\right]+(1-M[k\left]\right)·x_B\left[k\right]$

Here, $M\mathit{\left[}k\mathit{\right]}\mathit{\in }\mathit{\{}\mathit{0}\mathit{,}\mathit{1}\mathit{\}}$ is a randomly generated binary mask. AGATa performs augmentation by randomly selecting one of the three techniques above for each epoch, increasing both the diversity of the data and the generalization performance of the model.

3) AGATa’s performance and contributions

1. Improved prediction performance in classification and regression tasks:

Image 2. Performance Comparison of Various Existing Augmentation Techniques and the AGATa Augmentation Technique^[8]

Image 3. Performance Comparison between the Previous SoTA Algorithm and AGATa^[8]

AGATa outperformed existing models in terms of accuracy and mean square error (MSE), and in particular, outperformed various existing augmentation techniques as shown in Image 2. Furthermore, AGATa achieved higher performance than GBDT models such as XGBoost, which have shown strength on tabular data and outperformed recently proposed SoTA DNN algorithms such as T2G-Former and SAINT as shown in image 3.

Conclusion

Tabular data is utilized in many domains, including finance, healthcare, and manufacturing, and requires a flexible approach that can operate reliably regardless of real-time data processing or dataset size. To further advance tabular data research based on AGATa, LG AI Research's DI Lab aims to develop technologies that can be applied to a variety of environments. To do so, we are exploring ways to effectively handle various data environments that are frequently encountered in practice, such as noise, missing values, and unbalanced data.

We are also focusing on understanding the structural features of tabular data and augmentation techniques and developing an integrated framework that can be applied to various scenarios. These studies will contribute to combining data processing and training processes to produce practical results that can be applied in real-world industrial settings.

NeurIPS 2024 Series

EP. 1 [NeurIPS 2024 Series] Introduction to Mutual Information Estimation Benchmark Research
EP. 2 [NeurIPS 2024 Series] B-splines-Based Input Value Normalization Method for Tabular Data
EP. 4 [NeurIPS 2024 Series] A Study on Adaptive Information Routing Technology for Multimodal Time Series Forecasting