Technical Summary: Toward Privileged Foundation Models: LUPI for Accelerated and Improved Learning

Problem Statement

Training Tabular Foundation Models (TFMs) is computationally intensive and often slow to converge due to the combination of massive, heterogeneous datasets and high-dimensional parameter spaces. While TFMs leverage in-context learning (ICL) to generalize across tasks, they often struggle to efficiently reconstruct global dataset properties or infer latent data-generating processes solely from observed context. This results in prolonged pretraining times and suboptimal asymptotic generalization. The paper addresses the question of whether auxiliary information, available only during training, can systematically improve both learning speed and generalization for TFMs.

Methodology

The authors introduce PIQL (Privileged Information for Quick and Quality Learning), a framework grounded in the Learning Using Privileged Information (LUPI) paradigm. In this setting, a "teacher" provides auxiliary information during training that is unavailable at inference time. The methodology consists of three core components:

1. Construction of Privileged Information (PI)

The paper proposes two complementary forms of PI tailored to the synthetic-data regime of TFMs:

• Meta-Features (Weak PI): These are aggregate dataset-level statistics (e.g., first four moments and five quantiles) computed from the input data. Unlike strict LUPI, these are available at both training and inference. They serve to reduce the burden on the model's ICL capacity by explicitly providing global distributional properties, effectively freeing up representational depth.
• Generator Program Encoding (Strong PI): This form leverages the unique setting of TFMs, where training data is sampled from known synthetic priors (e.g., Gaussian Mixture Models, Structural Causal Models, Copulas). The "generator program" (the specific prior family and its parameters) is encoded into a sequence of tokens. This information is available only during training (train-time-only) and provides knowledge about the underlying data-generating process that cannot be inferred from finite samples alone.

2. Architecture and Transfer Mechanism

To utilize train-time-only PI (the generator program) during inference, the authors design a Teacher-Student TFM architecture:

• Teacher: During pretraining, the model receives the true generator program encoding.
• Student (Transfer Model): A lightweight model learns to reconstruct the generator program embedding from the observable context (meta-features and data points).
• Annealed Teacher Forcing: To mitigate the train-inference discrepancy (exposure bias), the training process employs scheduled sampling. The true generator tokens are progressively replaced by the student's predicted embeddings according to a decaying schedule, allowing the model to gradually learn to transfer the privileged signal.

3. Program Encoder

To handle heterogeneous priors (GMMs, SCMs, Copulas), a specialized Program Encoder is introduced. It serializes the generator specifications into a domain-specific language (DSL) format and tokenizes them. A Transformer encoder then maps these serialized programs into a fixed-size vector representation (prefix tokens) that can be prepended to the TFM context.

Key Contributions

• Novel Training Paradigm: PIQL is the first framework to systematically integrate privileged information into foundation model training to simultaneously accelerate convergence and improve generalization.
• Constructive PI Design: The paper moves beyond empirical trial-and-error by proposing two specific, theoretically motivated forms of PI: aggregate statistics to relieve ICL capacity and generator program encodings to expose latent generative structures.
• Architectural Innovation: A new architecture is designed to transfer train-time-only information via a student model trained with annealed teacher forcing, enabling the model to reconstruct privileged signals at inference.
• Theoretical Analysis: The authors provide theoretical proofs characterizing conditions under which PI reduces the population-level approximation gap. They distinguish between representation improvement (via meta-features bypassing architectural bottlenecks) and information improvement (via generator PI reducing uncertainty about latent structure).
• Empirical Validation: Extensive experiments demonstrate that PIQL enables TFMs to achieve faster convergence, lower final loss, and better generalization compared to baselines without PI.

Results

Empirical evidence from proof-of-concept studies and large-scale experiments (10-layer TFM on mixed priors) supports the theoretical claims:

• Accelerated Convergence: Privileged models reach lower training loss and higher test AUROC in significantly fewer epochs. For instance, PIQL models achieved comparable performance to baselines in roughly half the training time (e.g., 1.78x to 11.34x faster depending on the metric and PI grade).
• Improved Generalization: Models trained with PI achieve lower asymptotic loss and better generalization on out-of-distribution benchmarks (e.g., ADBench).
• Representation Efficiency: Layer-wise probing reveals that PIQL models capture latent distributional properties (such as generator type and cluster count) in earlier layers than non-privileged counterparts, indicating that PI effectively increases the model's "effective depth" by freeing ICL capacity.
• Ablation Insights:
  • Off-the-shelf LLM embeddings for program encoding failed to provide benefits, highlighting the necessity of the specialized program encoder.
  • Generator-PI (strong PI) provided more significant gains than Meta-PI (weak PI), though combining them yielded marginal improvements.
  • Expanding Meta-PI to 100 features (Meta-PI++) using LLM-generated statistics could match the performance of Silver-grade (generator) PI, suggesting a trade-off between feature richness and generative knowledge.

Significance and Claims

The paper positions PIQL as a principled and practical paradigm for improving the efficiency of foundation models. By formalizing the use of "teacher" information available only during training, the work demonstrates that:

1. Efficiency: Pretraining costs can be reduced by accelerating convergence, requiring fewer epochs and less compute.
2. Performance: Generalization capabilities can be enhanced by guiding the model with knowledge of the data-generating process that is otherwise inaccessible.
3. Paradigm Shift: The work establishes that leveraging the synthetic nature of TFM pretraining data (where the generator is known) is a viable and powerful source of privileged information, distinct from standard auxiliary data or prompt engineering approaches.

The authors conclude that this approach offers a pathway toward more efficient and effective foundation models, reducing the barriers to entry for developing such systems while maintaining or improving performance.