Technical Summary: Geometry-Aware Tabular Diffusion (GATD)

Problem Statement

Tabular data remains the dominant format in enterprise, healthcare, and scientific applications, yet synthesizing realistic tabular data for privacy-preserving sharing and augmentation presents unique challenges. Unlike images or text, tabular data exhibits heterogeneous column types, complex inter-column dependencies, and highly non-Gaussian marginal distributions. While diffusion models have emerged as a promising solution (e.g., TabDDPM, TabSyn, TabDiff), they typically rely on implicit mechanisms—such as self-attention in transformer architectures—to infer inter-column relationships. This paper posits that leaving relational structure to be inferred solely from the denoising objective is suboptimal and asks whether explicit relational structure can be provided as an auxiliary supervision signal that transfers across different denoising architectures.

Methodology

The authors introduce Geometry-Aware Tabular Diffusion (GATD), a framework that augments tabular diffusion denoisers with explicit pairwise geometric features computed directly from column values.

Geometric Feature Representation

The core innovation involves mapping all columns (continuous and categorical) to a unified normalized space $v \in [-1, 1]$. From this space, the model computes two pairwise features for every column pair $(i, j)$:

1. Pairwise Angle ($\theta_{ij}$): $\theta_{ij} = \arctan(v_j - v_i)$. This captures the directional relationship between columns, bounded within $(-\pi/2, \pi/2)$.
2. Pairwise Length ($\ell_{ij}$): $\ell_{ij} = \frac{1}{2} \log(1 + (v_j - v_i)^2)$. This captures the magnitude of the difference, with logarithmic compression for large differences.

Architecture and Training

GATD is implemented as a "drop-in" module across three distinct diffusion backbones:

• Residual MLP: The primary instantiation, extending a residual MLP with geometric inputs and auxiliary heads.
• GNN: Uses edge-conditioned message passing on a complete column graph, with geometric features as edge attributes.
• Transformer: Uses a column-wise transformer encoder, with geometric features projected and concatenated to the representation.

Input and Augmentation:
The model receives concatenated inputs including time embeddings, noised data, and the computed geometric features ($\theta_{input}, \ell_{input}$). Crucially, the architecture includes auxiliary prediction heads for angles and lengths.

• Training: The model is trained with a composite loss function:
  $$L = L_{cont} + L_{cat} + \lambda_\theta L_{angle} + \lambda_\ell L_{length} + \lambda_c L_{consistency}$$
  The geometric losses ($L_{angle}, L_{length}$) provide direct supervision. A consistency loss ($L_{consistency}$) ensures predicted geometry matches the denoised output. The authors note that geometric terms account for approximately 95% of the total loss at convergence, a deliberate inversion of typical loss hierarchies to force the network to internalize inter-column structure.
• Inference: During sampling, geometric features are computed from the current state and fed as inputs. The predicted angles ($\hat{\theta}$) are concatenated with the hidden state to form an augmented representation ($h_{aug}$) that feeds the denoising heads. The length head is detached during generation, serving only as a training regularizer.

Key Contributions

1. Explicit Geometric Supervision: The introduction of pairwise angle and length features as both inputs and auxiliary prediction targets for tabular diffusion.
2. Supervision vs. Capacity Ablation: An architecture-matched ablation demonstrating that the performance gain is driven by supervision, not merely the addition of geometric inputs or increased model capacity. An "InputsOnly" configuration (geometric features present but no geometric loss) performed indistinguishably from a baseline with no geometry ($d = -0.08$), while the supervised version showed a large effect ($d = 0.81$).
3. Cross-Architecture Portability: The geometric signal acts as a portable inductive bias, improving performance across MLP, GNN, and Transformer backbones without architecture-specific tuning.
4. Efficient MLP Instantiation: A compact MLP-based GATD that matches or exceeds the state-of-the-art TabDiff (Transformer-based) while using 3.5× fewer parameters on average (up to 25× fewer for classification tasks).
5. Practical Guidance: Empirical findings on loss weight tuning, sampling steps, and the handling of mixed data types.

Experimental Results

The method was evaluated on 10 benchmark datasets (5 classification, 5 regression) against baselines including TabDiff, TabDDPM, and STaSy.

• Fidelity: GATD (MLP instantiation) won 8/10 datasets on Shape (marginal distribution fidelity) and 7/10 on Trend (correlation preservation) compared to TabDiff. Across all three backbones (MLP, GNN, Transformer), +GEOM won 27/30 Shape cells and 25/30 Trend cells.
• Downstream Utility: The method reduced Shape error by 27% and Trend error by 20% on average compared to TabDiff. It closed 58.2% of the gap to real-data performance for AUROC/R2 and 42.0% for F1/RMSE.
• Efficiency: Despite using 2.5× more training epochs (20,000 vs. 8,000) and 20× more sampling steps (1000 vs. 50), the smaller MLP backbone resulted in a 1.7× faster end-to-end training time on average compared to TabDiff.
• Scalability: On a wide-table dataset (APS Failure, $d=171$), GATD maintained competitive performance, though memory usage scaled quadratically ($O(d^2)$), consistent with self-attention mechanisms.

Significance and Claims

The paper claims that explicit relational supervision is a portable inductive bias for tabular diffusion. The authors argue that while attention mechanisms can implicitly learn column relationships, providing explicit geometric targets reduces the burden on the denoiser and accelerates learning.

Key takeaways emphasized by the authors include:

• Supervision is Essential: The gain comes from the auxiliary prediction task forcing the network to internalize structure, not from the geometric inputs themselves.
• Complementarity: Geometric supervision and attention are complementary, not mutually exclusive; the signal improves transformer-based denoisers just as it does MLPs.
• Modesty on Privacy: The paper explicitly states that synthetic data generation provides no formal privacy guarantees. It warns against "privacy-washing" and encourages pairing synthetic generation with formal frameworks like differential privacy (DPSGD).
• Limitations: The method incurs $O(d^2)$ computational scaling due to pairwise features, which becomes a memory bottleneck for very wide tables, though it remains feasible on standard hardware for the evaluated regimes. The loss weights are heuristic and may require per-dataset tuning for optimal results in specific regimes (e.g., continuous-heavy datasets).