Prior-data Fitted Networks (PFNs) have reframed supervised tabular learning as single-pass in-context inference without per-dataset optimization. Extending this paradigm to clustering is appealing yet fundamentally more challenging, due to absent supervision, unknown cluster cardinality, and label switching inherent to partition outputs. Existing PFN-based clustering methods address these challenges only partially, either requiring known cardinality as input or relying on unstable label-ordering conventions and overly restrictive synthetic priors. We introduce TabClustPFN, a clustering PFN that resolves all challenges jointly through co-designed prior, objective, and architecture. Our hybrid pretraining prior captures heterogeneous real-tabular geometry; our decoupled partition inference network--cardinality inference network architecture jointly infers cluster assignments and cardinality in a single pass; and our SoftARI training objective is permutation-invariant by construction, eliminating the need for any label-ordering convention. On a 44 curated real-world tabular benchmark, TabClustPFN achieves state-of-the-art clustering performance against classical, deep, and amortized baselines, with runtime comparable to efficient baselines. Code is available at https://github.com/Tianqi-Zhao/TabClustPFN.
Figures & tables
Figure 1: Time & median ARI rank (lower is better). TabClustPFN achieves the best ARI with runtime comparable to simple methods and faster than other PFN-based methods.
Figure 2: TabClustPFN architecture & loss. Given dataset X and candidate cluster number K , the PIN uses a transformer encoder followed by L layers of iterative cross-attention between data representations and learnable cluster prototypes to produce soft assignments P . Dotted blue lines denote query paths, and dashed arrows indicate keys and values. Cosine similarity between normalized representations and prototypes yields cluster posteriors. For K∈{2,…,Kmax} , Gram matrices PK⊤PK are summarized and fed to an MLP-based CIN to estimate p(K∣X) . PIN is trained with permutation-invariant SoftARI loss, while CIN uses cross-entropy on true K∗ under decoupled optimization.
Method
ARI Rank [IQR] ↓
NMI Rank [IQR] ↓
k -MAE [IQR] ↓
TabClustPFN
2 [4]
3 [5]
0 [1]
AP
14 [1]
13 [3]
44 [12]
Agg +
9 [4]
10 [4]
1 [4]
Birch +
9 [3]
10 [3]
1 [4]
DBSCAN
16 [2]
17 [1]
5 [5]
GMM +
6 [4]
6 [5]
1 [3]
Table 1: Performance on synthetic datasets . Median ARI/NMI ranks and k -MAE ( ↓ ). Brackets show IQR. Bold and underline indicate best and second-best per column.
Figure 4
Nominal α
80%
85%
90%
95%
99%
Empirical coverage
0.875
0.907
0.937
0.968
0.993
Table 2: Calibration of the CIN posterior .
Figure 5: OOD performance under larger N (fixed D=50 ).
Figure 7: Benchmark performance on real-world benchmark . Median of ARI rank (left) and MAE in the estimated number of clusters (right); lower is better for both metrics. Red bar denotes our TabClustPFN , green bars denote other PFN-based methods, and yellow bars denote classical methods.
Method
ARI ↑
k -MAE ↓
k -Acc ↑
Time (ms) ↓
DAC (Anc)
0.5665
10.0
0.00625
257.20
DAC (Min)
0.5209
5.0
0.0375
235.48
NCP
0.6000
3.0
0.08125
1058.04
PFN (Random-init)
0.4222
1.0
0.4250
69.58
PFN (TabICL-init)
0.6294
1.0
0.4250
71.79
TabClustPFN
0.7893
1.0
0.4250
70.93
Table 3: Quantitative comparison of clustering performance and efficiency under the evaluation setting used by amortized clustering methods . Results are reported in terms of median ARI, k -MAE, k -Acc, and clustering time. Bold values indicate the best result for each metric ( ↑ higher is better, ↓ lower is better).
synthetic
Real
Method
ARI ↑
NMI ↑
ARI ↑
NMI ↑
Naïve (SA + MLP)
0.8137
0.7446
0.3589
0.3271
Non-iterative (Transformer)
0.8229
0.7523
0.3505
0.3865
Iterative (Ours)
0.8308
0.8000
0.3853
0.4237
Table 4: PIN decoder architecture ablation . Comparison of different model architectures on synthetic and real datasets using median ARI and NMI (higher is better). Bold and underlined values indicate the best and second-best performance in each column.
Method
Synthetic
Real
Time(ms) ↓
ARI ↑
NMI ↑
ARI ↑
NMI ↑
Matching+SoftAcc
0.8021
0.7364
0.3814
0.4071
70.0346
Matching+CE
0.8627
0.7637
0.3507
0.4070
6.2625
SoftNMI
0.7637
0.6929
0.3915
0.4250
3.3658
SoftARI (Ours)
0.8308
0.8000
0.3853
0.4237
2.9695
Table 5: PIN loss ablation . Median ARI and NMI on synthetic and real datasets.
Synthetic
Real
Method
ARI ↑
NMI ↑
ARI ↑
NMI ↑
Additive
0.8103
0.7227
0.3544
0.3716
Decoupled (Ours)
0.8308
0.8000
0.3853
0.4237
Table 6: Decoupled loss ablation . Median ARI and NMI on synthetic and real datasets.
Prior data
Synthetic ↑
Real ↑
GMM
ZEUS
GMM+ZEUS
GMM
0.8236
0.5434
0.7428
0.2270
ZEUS
0.6758
0.8047
0.7405
0.3414
GMM+ZEUS (Ours)
0.8230
0.8334
0.8308
0.3853
Table 7: Prior data types ablation . Ablation study of prior data types on synthetic and real datasets using median ARI.
Appendix figures & tables24 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: Visualization of randomly generated 2D datasets from the GMM prior .
Figure 9: Visualization of randomly generated 2D datasets from the ZEUS prior .
ID
# Instances
# Numerical features
# Categorical features
Dimension
# Classes
0
511
55
0
55
8
1
741
33
0
33
2
2
827
8
0
8
7
3
993
47
0
47
10
4
740
23
0
23
6
5
675
39
0
39
5
Appendix
Table 8: Statistics of synthetic datasets for evaluation generated with the GMM prior .
ID
# Instances
# Numerical features
# Categorical features
Dimension
# Classes
0
649
3
4
7
2
1
910
18
14
32
7
2
948
6
43
49
6
3
905
13
17
30
2
4
863
14
1
15
2
5
898
6
0
6
2
Appendix
Table 9: Statistics of synthetic datasets for evaluation generated with the ZEUS prior .
ID
# Instances
# Numerical features
# Categorical features
Dimension
# Classes
16
2000
64
0
64
10
18
2000
6
0
6
10
22
2000
47
0
47
10
51
294
6
7
13
2
53
270
13
0
13
2
481
209
7
1
8
2
Appendix
Table 10: Statistics of OpenML datasets used in the real-world benchmark .
Name
# Instances
# Numerical features
# Categorical features
Dimension
# Classes
haberman
306
2
1
3
2
banknote
1372
4
0
4
2
iris
150
4
0
4
3
thyroid
215
5
0
5
3
ecoli
336
7
0
7
8
seeds
210
7
0
7
3
Appendix
Table 11: Statistics of UCI datasets used in the real-world benchmark .
Method
ARI [IQR] ↑
NMI [IQR] ↑
k -MAE ↓
Low
High
Low
High
Low
High
TabClustPFN ∗
0.8348 [0.22]
0.1125 [0.10]
0.8012 [0.28]
0.1571 [0.14]
-
-
TabClustPFN
0.8360 [0.24]
0.1321 [0.08]
0.8046 [0.32]
0.1940 [0.10]
0
5
AP
0.0551 [0.10]
0.0243 [0.02]
0.3169 [0.28]
0.1440 [0.06]
44
45
Agg ∗
0.6586 [0.39]
0.1536 [0.07]
0.6513 [0.35]
0.1562 [0.01]
-
-
Agg +
0.6081 [0.46]
0.1278 [0.09]
0.6006 [0.40]
0.1087 [0.04]
1
3
Appendix
Table 12: Comprehensive performance on synthetic datasets by difficulty . Results show median [IQR] for ARI and NMI (higher is better), and k -MAE. † denotes method-specific selection criteria described in the text.
Method
ARI [IQR] ↑
NMI [IQR] ↑
k -MAE ↓
Low
High
Low
High
Low
High
TabClustPFN ∗
0.5754 [0.31]
0.1897 [0.21]
0.5737 [0.34]
0.1614 [0.13]
-
-
TabClustPFN
0.4431 [0.29]
0.1801 [0.14]
0.5311 [0.28]
0.1843 [0.16]
1
1
AP
0.1217 [0.18]
0.0298 [0.05]
0.3306 [0.29]
0.1408 [0.13]
27
40
Agg ∗
0.5104 [0.27]
0.1404 [0.06]
0.5665 [0.32]
0.0786 [0.05]
-
-
Agg +
0.5019 [0.37]
0.1359 [0.11]
0.4544 [0.34]
0.0915 [0.06]
0
2
Appendix
Table 13: Comprehensive performance on real-world benchmark by difficulty . Results show median [IQR] for ARI and NMI (higher is better), and k -MAE. † denotes method-specific selection criteria described in the text.
Method
ARI Rank [IQR] ↓
NMI Rank [IQR] ↓
TabClustPFN*
1 [4]
3 [5]
Agg*
8 [4]
9 [4]
Birch*
8 [5]
9 [4]
GMM*
6 [7]
5 [5]
KMeans*
6 [6]
6 [6]
Spectral*
13 [2]
13 [1]
Appendix
Table 14: Performance on synthetic datasets . Median ARI/NMI ranks of methods with known K . Brackets show IQR. Bold and underline indicate best and second-best per column.
Figure 10: Representation t -SNE visualizations on GMM prior . Rows are organized in groups of four; within each group, columns correspond to different datasets. From top to bottom, rows display raw features, TabICL encoder outputs, ZEUS representations, and TabClustPFN representations.
Figure 11: Representation t -SNE visualizations on ZEUS prior . Rows are organized in groups of four; within each group, columns correspond to different datasets. From top to bottom, rows display raw features, TabICL encoder outputs, ZEUS representations, and TabClustPFN representations.
Aspect
Cluster-PFN
TabClustPFN
Prior
Finite GMM prior in a low-dimensional setting ( D≤5 ).
Controlled-overlap GMM prior combined with the ZEUS prior, covering both Gaussian and nonlinear non-Gaussian cluster geometries ( D≤64 ).
Cluster-assignment objective
Uses deterministic relabeling to break label permutation symmetry, followed by cross-entropy.
Optimizes a permutation-invariant SoftARI loss, avoiding deterministic relabeling conventions that can be sensitive to scaling and dimensionality.
Cluster-assignment architecture
Uses standard Transformer layers followed by an MLP output head.
Uses a TabICL encoder followed by an iterative prototype–data cross-attention decoder in PIN.
Cluster-assignment mechanism
Outputs membership probabilities directly from point embeddings.
Maintains K cluster prototypes and computes assignments via cosine similarities between point representations and prototypes.
Cardinality-inference architecture
Constructs a special object ρ from a fixed all- −1 vector, concatenates it with dataset embeddings, and maps the decoded representation of ρ to p(K∣X) .
Uses CIN to map assignment Gram matrices {PK⊤PK}K=2Kmax to a posterior over K .
Conditioning on K
Conditions the model through an embedding of the candidate value of K .
Conditions PIN by selecting the first K cluster prototypes, so the active prototypes directly define the candidate cardinality.
Appendix
Table 15: Methodological comparison between Cluster-PFN and TabClustPFN . Both methods perform amortized clustering, but differ in their prior, objective, architecture, and cardinality-inference mechanism.
Figure 12: Comparison with Cluster-PFN under the low-dimensional GMM regime. All methods are evaluated in the unknown- K setting. Higher ARI/NMI and lower k -MAE indicate better performance. TabClustPFN consistently outperforms Cluster-PFN for D>2 .
Method
Synthetic ( 49 datasets)
Real-world ( 44 datasets)
ARI ↑
NMI ↑
k -MAE ↓
ARI ↑
NMI ↑
k -MAE ↓
TabClustPFN
0.8308 [0.39]
0.8000 [0.42]
0
0.3853 [0.42]
0.4237 [0.36]
1
DP-means
0.0552 [0.08]
0.3484 [0.29]
302
0.0351 [0.10]
0.2767 [0.25]
138
DP-means †
0.2892 [0.78]
0.4340 [0.67]
3
0.2774 [0.41]
0.3424 [0.36]
2
Dip-means
0.0000 [0.00]
0.0000 [0.00]
3
0.0000 [0.24]
0.0000 [0.44]
1
Dip-means †
0.0000 [0.00]
0.0000 [0.00]
4
0.0558 [0.39]
0.1245 [0.48]
1
Appendix
Table 16: Comparison with adaptive-cardinality baselines under unknown K . ARI and NMI are reported as median [IQR], and k -MAE as the median absolute cardinality error. Unmarked baselines use package defaults; † denotes silhouette-selected hyperparameters.
Figure 13: Paired t -test results based on per-dataset ARI differences Δb=ARITabClustPFN−ARIb . The dashed red line denotes zero difference. Positive values indicate that TabClustPFN outperforms the corresponding baseline. Colors indicate the one-sided paired t -test p -value under the alternative hypothesis H1:ARITabClustPFN>ARIb .
Figure 14: Wilcoxon signed-rank test results based on per-dataset ARI differences Δb=ARITabClustPFN−ARIb . Positive values indicate that TabClustPFN outperforms the corresponding baseline. Colors indicate the one-sided Wilcoxon signed-rank test p -value under the alternative hypothesis H1:ARITabClustPFN>ARIb .
Figure 15: Critical-difference (CD) diagrams based on ARI ranks. Lower average rank indicates better performance.
Figure 16: Pairwise win rate comparison on real-world benchmark . Each entry shows the percentage of datasets on which the row method achieves a higher ARI than the column method.
Figure 17: Runtime and median ARI rank grouped by knowledge of K on real-world benchmark . Top row: methods with known K . Bottom row: methods with unknown K .
Figure 18: Synthetic distribution shifts at D=2 . Left: cluster-size imbalance. Right: heavy-tailed subpopulations. Each point reports median ARI ten data-generation seeds ( 0 – 9 ) for one condition under unknown K .
Figure 19: Synthetic distribution shifts across feature dimensions. For each method and dimension, we first compute median ARI ten data-generation seeds ( 0 – 9 ) within each condition, then take the median across the seven imbalance ratios (left) or seven degrees of freedom (right).
OpenML ID
Dataset
N
D
K
458
analcatdata_authorship
841
70
4
12
mfeat-factors
2000
216
10
46980
MIC
1699
111
8
14
mfeat-fourier
2000
76
10
41143
jasmine
2984
144
2
40966
MiceProtein
1080
77
8
Appendix
Table 17: The 6 real-world OOD datasets. OpenML IDs identify the exact datasets. N and D are the original sample and feature counts.
Method
ARI ↑
NMI ↑
k -MAE ↓
Time (s) ↓
TabClustPFN
0.200 [0.371]
0.352 [0.462]
1
0.124
ZEUS+KMeans +
0.165 [0.441]
0.327 [0.487]
5
1.363
KMeans +
0.183 [0.469]
0.229 [0.523]
4
0.106
Agg +
0.155 [0.470]
0.289 [0.545]
4
1.792
HDBSCAN
0.083 [0.064]
0.141 [0.319]
7
1.016
Appendix
Table 18: Real-world OOD results on the 6 datasets in Table 17 . Median [IQR] ARI/NMI, median k -MAE, and mean clustering time. Bold marks the best result in each column.
CIN loss
Synthetic
Real
Order-aware
1.0
2.5
CE (Ours)
0.0
1.0
Appendix
Table 19: CIN loss ablation . Median k MAE on synthetic and real datasets using different CIN losses.
Prior-data fitted networks (PFNs) have achieved exceptional performance on tabular classification tasks. However, like other classifiers, their performance can suffer under the effect of class imbalance, resulting in poor performance for rare classes. Several techniques exist which attempt to mitigate the deleterious effect of class imbalance on classification performance, but the in-context learning (ICL) dynamic of PFNs means that loss-based strategies are impossible, and other techniques are unproven. We have adapted several classical techniques addressing class imbalance and analyzed their performance on PFN classification. We observe that thresholding performs exceptionally well because of the calibration characteristics of PFNs, and downsampling performs comparably because of PFNs exceptional limited-data performance, with the additional benefit of reduced computation cost for inference.
Samuel McDowell, Nathan Stromberg, Lalitha Sankar
School of Electrical, Computer and Energy Engineering · Arizona State University, Tempe, AZ
Prior-fitted networks (PFNs) are a promising class of tabular foundation models that perform in-context learning, whereby the entire labelled training set is supplied as context, and predictions for test queries are produced in a single forward pass. However, the quadratically scaling self-attention mechanism in many PFN architectures makes inference prohibitive for very large training datasets. We propose CRUMB (Clustered Retrieval Using Minimised-MMD Batching), a three-stage inference wrapper that (i) clusters the test queries, (ii) selects a small, distributionally matched training subset for each cluster by greedily minimising the maximum mean discrepancy (MMD), and (iii) runs exact PFN inference on each reduced-context batch. CRUMB is architecture-agnostic and requires no retraining. On the 51-dataset TabArena benchmark, evaluated across three PFN architectures (TabPFNv2, TabICLv1, TabICLv2), we show that CRUMB outperforms similar state-of-the-art context selection strategies. We also show that CRUMB is resilient to covariate drift, as the MMD-minimisation step naturally helps align the training context distribution to match the current test batch distributions.
Jamie Heredge, Mattia J. Villani, Pranav Deshpande +2
Global Technology Applied Research, JPMorganChase, New York, NY 10001, USA
Prior-Data Fitted networks (PFNs) have been very successful in tabular contexts, handling prediction tasks in context. However, they are designed for single-task inference, meaning that predicting several target values within a context requires repeated forward calls and precludes inter-task information sharing. We propose TabPFN-MT, which is trained on an expanded multi-target synthetic prior to capture inter-task dependencies in context. This model uses an expanded y-encoder and a shared decoder head to enable multitask in-context learning and simultaneous inference. The model is uniquely specialized for small-to-medium datasets by relying on in-context learning rather than traditional gradient-based training. Within this regime (averaging fewer than 1,000 samples), extensive evaluations across 344 datasets demonstrate that TabPFN-MT establishes a new state-of-the-art for deep tabular multitask learning. Furthermore, despite the inherent compute asymmetry of joint optimization, our model remains highly competitive with the latest state-of-the-art single-task ensembles. Notably, on multitask datasets it achieves an overall Accuracy rank of 4.89, the highest average rank among all models tested. Crucially, TabPFN-MT delivers this highly competitive performance while reducing the inference cost for T tasks from O(T) to O(1) forward passes, offering a massive computational efficiency improvement for multi-target tabular applications.
Cormac Cureton, Narges Armanfard
McGill University · Mila - Quebec AI Institute · Montreal, QC, Canada