🤖 machine learning

In-Context Density Estimation for Tabular Data

This paper introduces ICED, a transformer-based, in-context energy density estimator pretrained on a synthetic prior that performs density estimation, out-of-distribution detection, anomaly detection, and generative augmentation across diverse tabular datasets in a single forward pass without requiring retraining, hyperparameter tuning, or labels.

Original authors: Patryk Marszałek, Jacek Tabor, Marek Śmieja

Published 2026-08-11

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Patryk Marszałek, Jacek Tabor, Marek Śmieja

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to understand a new city. In the old days, if you wanted to know where the crowds usually gather, where the quiet parks are, or where the dangerous alleys hide, you had to hire a different team of cartographers for every single neighborhood. Each team would spend weeks studying just one block, drawing their own map, and then packing up to move to the next. It was slow, expensive, and required a different set of tools for every job.

In the world of computer science, specifically in a field called machine learning, this is exactly how we used to handle tabular data—the kind of data that looks like a spreadsheet with rows and columns. The goal is density estimation: figuring out where the "probability mass" lies. Think of this as a map of how crowded different areas are. If you know where the data is crowded, you can spot the weird outliers (anomalies), find things that don't belong (out-of-distribution detection), or even invent new, realistic fake data to help train other computers (data augmentation). For years, the rule was: "One dataset, one custom-trained model." But what if you could build one super-smart guide who learns the skill of mapping, rather than just memorizing one specific map?

This is where a new paper introduces ICED (In-Context Energy-based Density estimator). The researchers, from Jagiellonian University in Poland, have built a single, frozen AI model that acts like a universal city guide. Instead of retraining a new model for every new spreadsheet, ICED reads the data you give it as a "context" and instantly tells you how "dense" or crowded any specific point is. It does this without needing to learn, tune, or adjust anything for the new data. It's like having a guide who has studied millions of different cities during their training and can now walk into a brand-new city, look at the streets, and immediately point out the busy squares and the empty, suspicious corners in a single glance.

The Problem: The "One-Size-Fits-None" Approach

For a long time, if you wanted to find a needle in a haystack (an anomaly) or generate new hay (data augmentation), you had to build a custom machine for that specific haystack. You'd feed it data, tweak its knobs (hyperparameters), and hope it learned the shape of that particular pile. If you got a new haystack with a different shape, you had to start all over again.

The paper argues that this is inefficient. While other fields like language processing have moved toward "foundation models"—huge AI brains that learn general rules and can be applied to new tasks instantly—tabular data density estimation has been stuck in the old way. Even existing "foundation models" for tables usually only solve one specific problem, like classifying if an email is spam or predicting a house price. They don't give you the underlying map of where the data lives.

The Solution: A Universal Map-Maker

The authors created ICED, a model based on a Transformer (the same architecture behind many modern AI chatbots). Here is how it works, using a playful analogy:

Imagine you are training a detective. Instead of showing them one crime scene and asking them to solve it, you show them millions of different crime scenes, all with the answers written in the margins. You teach them to recognize patterns: "Oh, when the windows are broken and the floor is wet, the suspect is likely here."

ICED was trained on a synthetic prior. This means the researchers didn't just feed it real-world data; they built a massive, artificial universe of data. They created millions of fake datasets with all sorts of weird shapes: some were smooth and round (Gaussian), some had long, heavy tails (heavy-tailed), some were twisted into complex curves (flows), and some had mixed types of information (like numbers and categories). Crucially, because they made this data, they knew the exact "true density" of every single point.

The model learned to look at a dataset (the context) and a specific point (the query) and predict an energy. In this paper, "energy" is just a fancy word for "unnormalized density." Think of it like a height map: high energy means a crowded, safe, typical area; low energy means a lonely, suspicious, or empty area. The model doesn't need to calculate the exact total area of the city (which is mathematically impossible to do perfectly); it just needs to know which spots are higher than others.

What ICED Can Do (Without Breaking a Sweat)

Once ICED is trained, it is frozen. It never changes. You can then use this single, unchanging model to do four very different jobs, all at once, without retraining:

Density Estimation: It tells you how typical a data point is.
Anomaly Detection: It spots the weirdos. If a point has very low energy (it's in a "valley" far from the crowd), ICED flags it as an anomaly.
Out-of-Distribution (OOD) Detection: It tells you if a new piece of data doesn't belong to the group at all.
Data Augmentation: It can generate new, fake data points that look real. It does this by following the "energy gradient" (like a hiker walking uphill toward the most crowded areas) to create new samples that fit the pattern.

The Results: Fast, Accurate, and Universal

The researchers tested ICED against a huge list of other methods, from old-school statistics to modern deep learning models.

Speed: Because ICED doesn't need to train on the new data, it is incredibly fast. The paper notes it is roughly 50 times faster than some of the other advanced models (like TabPFN) because it does everything in a single forward pass. It's like taking a photo of a city versus spending weeks drawing it by hand.
Accuracy: On synthetic tests where they knew the ground truth, ICED was consistently the best or second-best at ranking points correctly. It didn't just guess; it understood the shape of the data.
Robustness: This is the most impressive part. When the "context" data was contaminated with noise (meaning the training data itself had some bad points mixed in), many other models crashed. Their performance dropped significantly. ICED, however, stayed in the top three. It was built to handle messy, heavy-tailed, and noisy data during its training, so it didn't panic when the real world got messy.
Data Generation: When used to create fake data to train other classifiers, ICED performed just as well as the best specialized tools (like TabEBM), often beating them.

The Catch (and the Future)

The paper is honest about its limits. ICED gives you a "relative" map. It can tell you that Point A is more crowded than Point B, but it doesn't give you a perfect, normalized probability number (like "there is a 45.2% chance this happens") without an extra step. Also, because it looks at the whole dataset at once, if your dataset is massive, it might get a bit slow, though the authors suggest this can be managed.

The authors don't claim to have solved every problem in the universe. They suggest that ICED is a powerful step toward treating density estimation as a reusable primitive. Just as we no longer build a new engine for every car, we might soon stop building a new density model for every spreadsheet. We can just use the one universal guide, ICED, to navigate any data landscape we throw at it.

In short, ICED is a "learn-once, use-forever" tool that turns the complex, slow process of mapping data into a single, instant glance. It proves that you don't need to reinvent the wheel for every new dataset; you just need a wheel that knows how to roll on any terrain.

Problem Statement
Density estimation is a foundational unsupervised task for tabular data, underpinning applications such as anomaly detection, out-of-distribution (OOD) detection, and data augmentation. Despite sharing a common probabilistic basis—determining where probability mass lies—these tasks are currently addressed in isolation. Practitioners typically fit a separate model for each dataset and task, requiring distinct hyperparameter tuning, optimization loops, and often re-engineering for specific objectives. This per-dataset paradigm contrasts sharply with the "foundation model" approach in natural language and computer vision, where large models acquire general capabilities during pretraining and adapt to new tasks via conditioning. While Prior-Fitted Networks (PFNs) have successfully extended this paradigm to supervised tabular learning (e.g., classification and regression), density estimation remains largely confined to per-dataset fitting methods like Kernel Density Estimation (KDE), Gaussian Mixture Models (GMM), and deep likelihood models. Existing in-context approaches for unsupervised tasks often produce task-specific outputs (e.g., cluster assignments or anomaly scores) rather than the underlying density, or they recover density inefficiently as a by-product of supervised models.

Methodology: ICED
The authors introduce ICED (In-Context Energy-based Density estimator), a transformer-based model designed to learn the density-estimation algorithm itself rather than fitting a model to every new dataset.

Architecture and In-Context Learning: ICED follows the PFN paradigm. It is a transformer pretrained once on a synthetic prior. During inference, it receives an unseen dataset $D$ as context and a set of query points $Q$ . It predicts the energy (unnormalized log-density) for each query point in a single forward pass, without gradient updates, sampling, or hyperparameter selection.
Energy Formulation: Instead of predicting a normalized probability density, ICED predicts an energy function $f_\theta(q | D) \approx \log p(q) + c_D$ , where $c_D$ is a dataset-dependent additive constant. This formulation bypasses the intractable partition function while retaining the relative ordering of points, which is sufficient for ranking, detection, and gradient-based generation.
Synthetic Pretraining Prior: The model is trained on a diverse synthetic prior constructed to mimic real-world tabular complexities while providing exact ground-truth log-densities. The prior includes:
- Base Mixtures: Gaussian and heavy-tailed (Student-t, Laplace, Cauchy) mixtures with random parameters.
- Nonlinear Warping: Invertible normalizing flows (RealNVP and a custom piecewise flow) applied to base mixtures to create curved, multimodal supports.
- Mixed Types: Randomly appended categorical features to handle mixed-type data.
- Perturbations: Query points are perturbed (via Gaussian noise, Mixup, or CutMix) to regularize the energy landscape in out-of-distribution regions.
Training Objective: The loss function combines two terms to address the specific challenges of density estimation:
1. Regression Loss ( $L_{reg}$ ): A squared error for queries in high-density (in-distribution) regions and a one-sided penalty for low-density (OOD) queries to bound the tail without forcing exact values on outliers.
2. Ranking Loss ( $L_{rank}$ ): A pairwise margin-based loss that penalizes incorrect ordering of query points. This ensures the model captures the relative density structure, which is critical for anomaly and OOD detection even if absolute energy values are imperfect.

Key Contributions

First In-Context Foundation Model for Tabular Density: ICED extends the PFN paradigm from supervised prediction to unsupervised probabilistic modeling, learning the density-estimation algorithm itself.
Specialized Synthetic Prior and Objective: The authors design a pretraining prior spanning multimodal, heavy-tailed, and mixed-type distributions, coupled with a hybrid regression-ranking loss that focuses model capacity on informative regions while preserving global ordering.
Unified Multi-Task Capability: A single frozen ICED model drives four traditionally separate tasks: density estimation, OOD detection, unsupervised anomaly detection, and generative data augmentation, requiring no retraining or task-specific adaptation.

Experimental Results
The authors evaluate ICED across four domains, comparing it against classical estimators (KDE, GMM), deep density models (MAF, RealNVP, VAE), and other tabular foundation models (TabPFN, TabICL, TabEBM).

Density Estimation: On held-out synthetic distributions, ICED achieves competitive or superior pairwise ordering accuracy compared to specialized methods. Unlike autoregressive in-context models (e.g., TabPFN) which are slow due to sequential feature generation, ICED produces joint densities in a single forward pass, making it roughly 50x faster than TabPFN and comparable to the fastest baselines.
Out-of-Distribution (OOD) Detection: On the ADBench benchmark with clean contexts, ICED achieves a mean AUROC of 88.3, ranking third among 26 methods and outperforming specialized deep density estimators and TabPFN.
Unsupervised Anomaly Detection: In settings where the context is contaminated with anomalies, ICED maintains robustness, achieving a mean AUROC of 76.7. It remains in the top three across both clean and contaminated settings, whereas methods that excel on clean data (like GMM and TabICL) suffer significant performance drops when the context is noisy.
Data Augmentation: Using Stochastic Gradient Langevin Dynamics (SGLD) on the learned energy surface, ICED generates synthetic data that improves downstream classification accuracy. It matches the performance of TabEBM (72.8% vs 73.0% accuracy) and outperforms the vanilla baseline and other generative augmenters, despite using the same frozen model for all other tasks.

Significance and Claims
The paper positions ICED as a step toward treating density estimation as a reusable, amortized primitive for tabular data, analogous to how PFNs have reshaped supervised learning. The primary significance lies in the elimination of the per-dataset optimization loop. ICED demonstrates that a single, frozen model can effectively handle diverse unsupervised tasks without retraining, tuning, or labels. The authors claim that by modeling an unnormalized energy and training with an objective that balances local accuracy with global ordering, they achieve a level of versatility and efficiency previously unattainable in tabular density estimation, bridging the gap between foundation models and unsupervised probabilistic tasks.

The Problem: The "One-Size-Fits-None" Approach

The Solution: A Universal Map-Maker

What ICED Can Do (Without Breaking a Sweat)

The Results: Fast, Accurate, and Universal

The Catch (and the Future)

More like this