Transformers for Tabular Data: A Training Perspective of Self-Attention via Optimal Transport

TL;DR

This work reframes Transformer self-attention for tabular data through Optimal Transport, tracking how attention remaps inputs during training and measuring its proximity to OT optima. It demonstrates that, while the final SA mapping often approaches OT-optimal coupling, the training path is inefficient, prompting two remedies: pretraining the MLP and an OT-based remapping via class-specific dummy Gaussians. The OT-based approach achieves comparable accuracy to Transformers with lower computational cost and better scaling under standardized inputs, though it depends on careful dummy-geometry design. Applications to synthetic and real data (Bangalore EEG) illustrate the potential for efficient, geometry-aware remappings in tabular tasks, at the cost of reduced transferability and task-specific design considerations.

Abstract

This thesis examines self-attention training through the lens of Optimal Transport (OT) and develops an OT-based alternative for tabular classification. The study tracks intermediate projections of the self-attention layer during training and evaluates their evolution using discrete OT metrics, including Wasserstein distance, Monge gap, optimality, and efficiency. Experiments are conducted on classification tasks with two and three classes, as well as on a biomedical dataset. Results indicate that the final self-attention mapping often approximates the OT optimal coupling, yet the training trajectory remains inefficient. Pretraining the MLP section on synthetic data partially improves convergence but is sensitive to their initialization. To address these limitations, an OT-based algorithm is introduced: it generates class-specific dummy Gaussian distributions, computes an OT alignment with the data, and trains an MLP to generalize this mapping. The method achieves accuracy comparable to Transformers while reducing computational cost and scaling more efficiently under standardized inputs, though its performance depends on careful dummy-geometry design. All experiments and implementations are conducted in R.

Figures (9)

• Figure 1: Transformer architecture (adapted from vaswani2017attention). The decoder and other components not relevant for tabular data are grayed out, as this work focuses exclusively on the encoder.
• Figure 2: Illustration of the difference between the Monge formulation (left) and the Kantorovich relaxation (right). In the Monge case, when both measures have the same number of points and equal weights, the optimal plan corresponds to a permutation matrix that pairs each $x_i$ with exactly one $y_j$. In contrast, the Kantorovich formulation allows for probabilistic couplings between arbitrary discrete measures that cannot be matched one-to-one. Black segments indicate the nonzero entries of the optimal coupling matrix $P$. Figure adapted from Computational Optimal Transport by Gabriel Peyré and Marco Cuturi.
• Figure 3: Example of learned remapping of the SA after 45 iterations. On the left the input data (scaled), on the right the learned remapping.
• Figure 4: Transformer-induced transport: on the left, the SA matching between input points (empty) and their final remapped positions (full); on the right, the full Transformer path showing the sequence of projections through training.
• Figure 5: Wasserstein distances for two-class point clouds with increasing separation (2, 4, 6, 8), corresponding to decreasing classification difficulty.