🤖 machine learning

CRUMB: Efficient Prior Fitted Network Inference via Distributionally Matched Context Batching

This paper introduces CRUMB, an architecture-agnostic inference wrapper that significantly improves the efficiency and performance of Prior-Fitted Networks by clustering test queries and selecting distributionally matched training subsets via MMD minimization, thereby enabling effective in-context learning on large datasets without retraining.

Original authors: Jamie Heredge, Mattia J. Villani, Pranav Deshpande, Akshay Seshadri, Niraj Kumar

Published 2026-06-11

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jamie Heredge, Mattia J. Villani, Pranav Deshpande, Akshay Seshadri, Niraj Kumar

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

The Problem: The Overwhelmed Librarian

Imagine you have a brilliant librarian (this is the Prior-Fitted Network, or PFN). This librarian has read every book in the world during their training. If you ask them a question, they don’t need to study your specific topic; they just use what they already know to give you an answer instantly.

However, there is a catch. To answer your question, the librarian needs you to hand them a stack of reference books (the training data) that are relevant to your question.

The Old Way: If you have 10,000 reference books, the librarian has to read all of them to find the answer. This takes a huge amount of time and brainpower (memory). It’s like asking the librarian to read an entire library just to answer one question.
The "Quick" Way: You could just pick 10 random books to hand them. This is fast, but those 10 books might be about cooking when your question is about physics. The answer will be bad.
The "Personal" Way: For every single question you ask, you find the 10 most relevant books. This gives a great answer, but it’s incredibly slow. If you have 1,000 questions, the librarian has to start reading from scratch 1,000 times.

The Solution: CRUMB (Clustered Retrieval Using Minimised-MMD Batching)

The authors propose a method called CRUMB. Think of CRUMB as a smart assistant who organizes the work so the librarian doesn’t get overwhelmed, but still gets the right books.

CRUMB works in three stages:

Stage 1: Grouping the Questions (Clustering)

Instead of handing the librarian 1,000 individual questions one by one, CRUMB looks at all the questions and groups them into 20 buckets (clusters).

Analogy: Imagine you have a pile of 1,000 mixed-up mail letters. Instead of sorting them one by one, you quickly sort them into 20 piles: "Bills," "Magazines," "Family," "Work," etc. Now you have 20 distinct topics instead of 1,000 individual items.

Stage 2: Finding the Right Books for Each Bucket (MMD Selection)

For each of the 20 buckets of questions, CRUMB goes to the library (the training data) and picks out a small, specific set of reference books.

The Secret Sauce (MMD): CRUMB doesn’t just pick books that are close to the questions (like finding the nearest neighbor). It uses a mathematical trick called Maximum Mean Discrepancy (MMD).
Analogy: Imagine the "Bills" bucket contains questions about electricity, water, and internet. A simple method might just pick books about electricity because it’s the most common. But CRUMB uses MMD to ensure the selected books cover the whole variety of the bucket. It picks a few books on electricity, a few on water, and a few on internet, ensuring the "distribution" of the books matches the "distribution" of the questions. It avoids picking 10 books on electricity and zero on water.

Stage 3: The Librarian Works in Batches (Batched Inference)

Now, the librarian takes the first bucket (e.g., "Bills") and the specific set of books selected for it. They read those books and answer all the questions in that bucket at once. Then they move to the next bucket.

Result: Instead of 1,000 separate reading sessions, the librarian only does 20 reading sessions. This is much faster.

Why is CRUMB Better?

The paper compares CRUMB to other methods and finds two main advantages:

Speed vs. Accuracy Balance:
- Uniform Subsampling (picking random books) is fast but inaccurate.
- kNN (picking the nearest book for every single question) is accurate but very slow.
- CRUMB is almost as accurate as kNN but much faster because it batches the work. It achieves this by ensuring the books it picks are statistically representative of the questions in that batch.
Resilience to "Drift" (Changing Trends):
- Imagine the questions you ask change over time. Maybe last month you asked mostly about winter coats, but this month you’re asking about swimsuits. This is called Covariate Drift.
- Other methods (like MICP) organize their books based on how the library is arranged. If the library is organized for winter, it struggles when you suddenly ask about summer.
- CRUMB organizes based on your current questions. Because it looks at your current batch of questions first, it naturally adapts. If your questions shift to swimsuits, CRUMB immediately selects books about swimsuits. The paper shows that as the "drift" gets worse, CRUMB stays much more accurate than other methods.

Summary

CRUMB is a wrapper that sits in front of a powerful AI model. It says:

"Don't look at every question individually."
"Group similar questions together."
"For each group, pick a diverse set of reference data that matches the group's needs."
"Answer the whole group at once."

This allows the AI to handle massive amounts of data quickly without losing accuracy, and it handles changing data trends better than previous methods.

Technical Summary: CRUMB – Efficient Prior Fitted Network Inference via Distributionally Matched Context Batching

Problem Statement
Prior-fitted networks (PFNs), such as TabPFN and TabICL, are tabular foundation models that perform in-context learning by treating the entire labeled training set as context for a single forward pass. While these models have demonstrated strong performance on small-to-medium datasets, often outperforming gradient boosting methods like CatBoost and XGBoost, their utility is severely limited by computational scaling. The self-attention mechanism in PFNs scales quadratically with the size of the training dataset ( $O(N^2)$ ), making inference prohibitive for large datasets. Existing solutions involve architectural changes (e.g., linear attention) or fine-tuning, which this work does not pursue. A natural alternative is context selection—replacing the full training set with a smaller subset. However, uniform subsampling discards information, while query-dependent methods like k-nearest neighbours (kNN) prevent efficient batching because each test point requires a unique context and separate forward pass.

Methodology
The authors propose CRUMB (Clustered Retrieval Using Minimised-MMD Batching), a three-stage inference wrapper that is architecture-agnostic and requires no model retraining. The method balances context quality with batching efficiency:

Test Query Clustering: The test set $D_{test}$ is partitioned into $K$ clusters $\{C_1, \dots, C_K\}$ using k-means on standardized input features. This groups test queries with similar characteristics, allowing them to share a common training context.
Distributionally Matched Context Selection: For each test cluster $C_k$ , a training subset $S_k$ of size $n \ll N$ is selected from the full training pool $D_{train}$ . The selection criterion is the minimization of the Maximum Mean Discrepancy (MMD) between the empirical distribution of the test cluster and the selected training subset. This is solved via greedy kernel herding. The MMD objective encourages diversity within the selected subset (penalizing redundancy) while rewarding proximity to the test cluster (ensuring relevance). The authors employ Random Fourier Features (RFF) and batched greedy selection to accelerate this process.
Batched PFN Inference: The PFN runs $K$ independent forward passes. For each cluster $k$ , the model receives the shared context $S_k$ and the batch of test points $C_k$ . This reduces the attention cost from $O(T \cdot N)$ to $O(T \cdot n)$ and reduces the number of forward passes from $T$ (per-query) to $K$ (per-cluster).

Key Contributions

MMD-Based Context Retrieval: The paper introduces greedy MMD minimization (kernel herding) as a mechanism for selecting training subsets. It demonstrates that this approach outperforms centroid-nearest-neighbour and Voronoi-uniform retrieval within the same clustering framework.
Batched Context Selection via Test-Side Clustering: By clustering test queries and sharing a single training context per cluster, CRUMB enables efficient batching while preserving query-relevant context, resolving the tension between kNN’s accuracy and its inability to batch.
Resilience to Covariate Drift: The MMD-minimization step naturally aligns the training context distribution with the current test batch distribution. The authors show that this provides resilience against covariate drift, where test inputs are drawn from a shifted distribution $P_{test}(x) \neq P_{train}(x)$ .

Results
The method was evaluated on the 51-dataset TabArena benchmark across three PFN architectures (TabPFNv2, TabICLv1, TabICLv2).

Performance vs. Baselines: On TabArena, CRUMB significantly outperformed uniform subsampling and the Mixture of In-Context Prompters (MICP) baseline at identical context budgets. CRUMB’s performance was statistically close to per-query kNN (which serves as an upper bound for context selection accuracy) but required only $K=20$ forward passes compared to $T=200$ for kNN, resulting in a ~10x reduction in inference time (35s vs. 361s).
Consistency: CRUMB’s advantage held across different PFN backbones, sampling proportions ( $n = 0.05N$ to $0.8N$ ), and training set sizes.
Large Datasets: On large datasets ( $N > 50,000$ ), CRUMB with early stopping (which adapts context size based on MMD convergence) won more often than MICP, particularly for TabPFNv2 and TabICLv1.
Covariate Drift: In controlled experiments simulating covariate drift, CRUMB’s advantage over MICP grew monotonically with drift intensity. At no drift ( $\tau=0$ ), CRUMB led by +4.9%; at complete out-of-sample drift ( $\tau=1$ ), the gap widened to +17.1%. MICP collapsed under drift because its training-side clusters became misaligned with the shifting test distribution, whereas CRUMB’s test-side clustering adapted automatically.

Significance and Limitations
The paper claims that CRUMB makes PFN inference practical for large tabular datasets by combining test-side clustering with distributional context retrieval. The significance lies in the synergy between clustering (which creates localized query groups) and MMD herding (which ensures distributional alignment rather than mere geometric proximity).

The authors note several limitations:

CRUMB requires the full set of test queries upfront for clustering and does not natively support online/streaming settings without modification (which reduces it to a MICP-like procedure).
The method relies on Euclidean k-means and an isotropic Gaussian RBF kernel, which may be poorly suited for high-dimensional or mixed-type feature spaces.
As an inference wrapper, CRUMB inherits the performance ceiling of the underlying PFN; if the model’s pretraining priors are a poor match for the data, context selection cannot fully compensate.