In-Context Density Estimation for Tabular Data

ICED is an in-context, energy-based density estimator for tabular data that removes the per-dataset cost of traditional density estimation. It is a transformer-based model pretrained once on a synthetic prior built specifically for density estimation, under an objective that fits log-density where it is informative and preserves its ordering elsewhere. At inference, it reads a dataset as context and returns an unnormalized log-density (an energy) for any query point in a single forward pass, with no fitting, sampling, or hyperparameter selection. A single frozen ICED model then drives four tasks usually handled by four specialized pipelines: density estimation, out-of-distribution detection, unsupervised anomaly detection, and generative augmentation. Across all four, it is competitive with the strongest task-specific method, while being the only approach that needs no retraining, no tuning, and no labels to move between them.

The paper introduces the first in-context foundation model for tabular density estimation, extending the PFN paradigm from supervised prediction to unsupervised probabilistic modeling. It designs "a synthetic pretraining prior and a training objective tailored to density estimation, combining exact synthetic log-densities, broad distributional diversity, and a hybrid regression-ranking loss that concentrates model capacity on the informative regions of the density landscape. It demonstrates that a single frozen pretrained model is competitive with specialized methods across four traditionally separate tasks (density estimation, OOD detection, anomaly detection, and data augmentation) while requiring no retraining, no hyperparameter search, and no task-specific adaptation."

The method formalizes in-context density estimation as learning a single model fθ that approximates Ep∼p(D)[log p(q) D], so that density estimation on a previously unseen dataset reduces to a single forward pass. ICED estimates an energy, i.e., an unnormalized log-density, rather than a normalized one, with its output approximating log p(q) only up to a dataset-dependent additive constant cD that absorbs the log-partition function. This choice is deliberate because the tasks that consume tabular densities, such as anomaly and OOD scoring, and gradient-based generation, are invariant to cD, depending only on relative density between points.

The synthetic prior is designed to span the distributional structures encountered in real tabular data while keeping exact log-densities tractable. Each training instance is generated by sampling a base mixture (Gaussian or heavy-tailed with K ∈ [2, 20] components), optionally warping it through an invertible flow (RealNVP or a piecewise flow), optionally appending one to five discrete features, and perturbing a fraction of the query points. The prior includes heavy-tailed mixtures to teach the model that density can decay at very different rates, which is essential for robust anomaly scoring. Nonlinear warping via flows injects nonlinear geometry while preserving an exact likelihood. Discrete features reflect the mixed-type nature of tabular data, with categorical columns one-hot encoded and interleaved with continuous columns in a random order. Perturbation for a robust landscape shapes the learned energy surface in out-of-distribution regions by perturbing a random 10% of the coordinates of a fraction of the query points using Gaussian noise, Mixup, or CutMix feature swaps. Targets are normalized per-instance using empirical 10th and 90th percentiles as anchors.

The training objective combines a regression term and a ranking term: L = Lreg + Lrank. For queries whose target lies in the in-distribution range, a squared error is used; for low-density queries, only a one-sided penalty ensures their predicted energy stays below the range. The ranking term penalizes every pair whose predicted order disagrees with the ground truth, using a logistic loss. This improves rank metrics (Kendall's τ, pairwise accuracy) even when absolute energies are imperfect, which is precisely what anomaly and OOD scoring rely on. The model is optimized with AdamW under a cosine schedule, with gradient accumulation and clipping, on a stream of synthetic instances generated on the fly. During pretraining, the context length varies from 200 to 2000 samples, the query length is fixed at 256, and the number of features ranges from 2 to 50.

The model architecture is a transformer that processes context and query points jointly as a set, in the spirit of TabPFN. Each point is embedded into a token by the TabICLv2 feature extractor, which handles inputs of varying dimensionality. The feature extractor is initialized with weights from a pretrained TabICLv2 model and fine-tuned jointly. Context and query tokens are processed together by a stack of transformer encoder layers under an attention mask in which query tokens attend to the full context while context tokens attend only among themselves. A small multilayer-perceptron head maps each query token to a scalar energy, so that all estimates are produced in one forward pass. The model is permutation-invariant in the data points.

At inference, the user supplies a context dataset D and queries Q, both scaled using statistics from D, and the frozen model returns energies in a single forward pass. The same model and energies serve every downstream task: (1) density estimation and ranking use energies directly as log-density proxies; (2) anomaly/OOD detection uses the negated energy as an outlier score; (3) data augmentation draws new in-distribution samples by following the energy gradient with Stochastic Gradient Langevin Dynamics (SGLD), conditioning on a class-specific context as in TabEBM.

In the density estimation experiment on held-out synthetic data, ICED is evaluated by rank agreement (pairwise ordering accuracy and Kendall's τ) between predicted and ground-truth log-densities. The pattern is not a single winner but a robustness gap: ICED is the best or second-best estimator in every distribution class, while the baselines are each strong only where their assumptions hold. ICED is also among the fastest methods, roughly 50× faster than TabPFN while producing a joint density directly.

In out-of-distribution detection on the ADBench suite, ICED obtains 88.3 mean AUROC, placing third among all 26 methods, behind only TabICL (89.2) and TACTIC-C (88.7), and well ahead of TabPFN (84.6) and TabEBM (52.9). In unsupervised anomaly detection with contaminated context, ICED reaches 76.7 mean AUROC, again third overall, behind TACTIC-N (78.7) and MAF (77.3). ICED is the one method that stays in the top three across both settings, attributed to the prior spanning heavy-tailed and perturbed distributions and the ranking term in the loss.

In data augmentation, following the TabEBM protocol, ICED reports 72.8 downstream classification accuracy, essentially matching the purpose-built TabEBM (73.0) and improving over the vanilla baseline (71.0), while reusing the same frozen model. Both energy-based methods clearly outperform the dedicated generative augmenters.

The ablation study removes one ingredient of the prior at a time. On the hardest setting (evaluation data with both categorical features and perturbations), the full prior is best (0.7209 pairwise accuracy), and it is also strongest or second-strongest on the aggregate rows. The loss ablation shows the full loss is best on every evaluation class; dropping the ranking term costs a consistent margin, and removing the regression term is more damaging still.

Limitations include: the accuracy of ICED's density estimation depends on how well the synthetic prior reflects the target distribution, with no formal coverage guarantee; ICED returns an unnormalized energy, which cannot serve as an absolute normalized probability without an extra normalization step; inference cost grows with context size; and the augmentation results use SGLD sampling, which reintroduces an iterative procedure at generation time even though density estimation remains a single forward pass.

Improvements to AI systems:

1. Unified in-context density estimation for tabular data: Replace task-specific pipelines (density estimation, OOD detection, anomaly detection, augmentation) with a single frozen transformer that reads any tabular dataset as context and outputs an energy (unnormalized log-density) for any query in one forward pass. This eliminates per-dataset training, hyperparameter tuning, and labeled data requirements across all four tasks.

2. Robust anomaly and OOD scoring via heavy-tailed and perturbed priors: Train the model on synthetic priors that include heavy-tailed mixtures and perturbed query points (Gaussian noise, Mixup, CutMix). This makes the learned energy landscape decay at realistic rates and remain informative in low-density regions, improving outlier detection on contaminated or unseen data without retraining.

3. Hybrid regression-ranking loss for density ordering: Optimize a combined loss that fits absolute log-density values only where informative and enforces pairwise ordering everywhere. This yields high rank accuracy (Kendall’s τ, pairwise accuracy) even when absolute energies are imperfect, directly benefiting tasks that rely on relative density (anomaly scoring, OOD detection).

4. Single-pass generative augmentation via energy gradients: Use the frozen model’s energy as a score function for Stochastic Gradient Langevin Dynamics (SGLD) to draw new in-distribution samples conditioned on class-specific context. This matches purpose-built generative augmenters (e.g., TabEBM) in downstream classification accuracy while reusing the same model—no separate generative training needed.

5. Permutation-invariant, variable-dimensionality handling: Process context and query points jointly as a set with a TabICLv2 feature extractor that handles 2–50 features of mixed continuous/discrete types. This enables zero-shot generalization to new datasets with different schemas and feature counts without architectural changes.

6. Fast, scalable inference for large contexts: Leverage the transformer’s set-based attention (queries attend to full context; context tokens attend only among themselves) to produce all query energies in one forward pass, achieving 50× speedup over prior in-context methods (e.g., TabPFN) while directly outputting joint densities.

7. Calibration-free, normalization-invariant outputs: Return energies that are invariant to dataset-dependent additive constants (absorbing log-partition function). This makes the model directly usable for ranking, thresholding, and gradient-based tasks without needing to compute normalized probabilities.

What the improved AI system can do:

• Given any unlabeled tabular dataset (e.g., medical records, sensor logs, financial transactions), it can instantly: (a) rank points by density, (b) flag outliers or out-of-distribution samples, (c) generate synthetic in-distribution data for class-balanced augmentation—all with a single frozen model and no retraining, no hyperparameter search, and no labels.

• It can switch between these tasks on-the-fly by simply changing the query set or applying SGLD, making it a drop-in replacement for four separate specialized pipelines in production environments.

• It remains competitive with or superior to task-specific baselines (e.g., 88.3 AUROC for OOD, 76.7 for anomaly detection, 72.8% augmentation accuracy) while being the only method that achieves top-three performance across both OOD and anomaly settings simultaneously.

• It handles mixed-type, high-dimensional, and nonlinearly warped data distributions without assumptions, thanks to the diverse synthetic prior (Gaussian/heavy-tailed mixtures, RealNVP/piecewise flows, discrete features, perturbations).

Density estimation underlies many unsupervised tasks on tabular data such as anomaly detection, out-of-distribution detection, and data augmentation. Although all these problems reduce to questions about where probability mass lies, they are typically solved individually by fitting a separate model to each dataset, with its own hyperparameters and tuning budget. We introduce ICED, an in-context, energy-based density estimator that removes this per-dataset cost. ICED is a transformer-based model pretrained once on a synthetic prior built specifically for density estimation under an objective that fits log-density where it is informative and preserves its ordering elsewhere. In the inference, it reads a dataset as context and returns an unnormalized log-density for any query point in a single forward pass, with no fitting, sampling, or hyperparameter selection. A single frozen ICED model then drives four tasks usually handled by four specialized pipelines: density estimation, out-of-distribution detection, unsupervised anomaly detection, and generative augmentation. Across all four, it is competitive with the strongest task-specific method, while being the only approach that needs no retraining, no tuning, and no labels to move between them. The code is available at https://github.com/gmum/iced.