TabClustPFN: A Prior-data Fitted Network for Tabular Data Clustering

Tianqi Zhao; Guanyang Wang; Yan Shuo Tan; Qiong Zhang

cs.LGJan 29, 2026

TabClustPFN: A Prior-data Fitted Network for Tabular Data Clustering

Authors: Tianqi Zhao, Guanyang Wang, Yan Shuo Tan, Qiong Zhang

Organizations: Renmin University of China · Rutgers University · National University of Singapore

Abstract

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

{\bm{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\in\{2,\ldots,K_{\max}\}

, Gram matrices

P_{K}^{\top}P_{K}

are summarized and fed to an MLP-based CIN to estimate

p(K|{\bm{X}})

. PIN is trained with permutation-invariant SoftARI loss, while CIN uses cross-entropy on true

K^{*}

under decoupled optimization.

Method	ARI Rank [IQR] $\downarrow$	NMI Rank [IQR] $\downarrow$	$k$ -MAE [IQR] $\downarrow$
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 (

\downarrow

). Brackets show IQR. Bold and underline indicate best and second-best per column.

Nominal $\alpha$	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 $\uparrow$	$k$ -MAE $\downarrow$	$k$ -Acc $\uparrow$	Time (ms) $\downarrow$
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 (

\uparrow

higher is better,

\downarrow

lower is better).

	synthetic		Real
Method	ARI $\uparrow$	NMI $\uparrow$	ARI $\uparrow$	NMI $\uparrow$
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) $\downarrow$
Method	ARI $\uparrow$	NMI $\uparrow$	ARI $\uparrow$	NMI $\uparrow$	Time(ms) $\downarrow$
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 $\uparrow$	NMI $\uparrow$	ARI $\uparrow$	NMI $\uparrow$
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 $\uparrow$			Real $\uparrow$
Prior data	GMM	ZEUS	GMM+ZEUS	Real $\uparrow$
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

Panel of Figure 8 — Figure 8: Visualization of randomly generated 2D datasets from the GMM prior .

Panel of Figure 9 — 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] $\uparrow$		NMI [IQR] $\uparrow$		$k$ -MAE $\downarrow$
	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] $\uparrow$		NMI [IQR] $\uparrow$		$k$ -MAE $\downarrow$
	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] $\downarrow$	NMI Rank [IQR] $\downarrow$
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.

Panel of Figure 10 — 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.

Panel of Figure 11 — 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\leq 5$ ).	Controlled-overlap GMM prior combined with the ZEUS prior, covering both Gaussian and nonlinear non-Gaussian cluster geometries ( $D\leq 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 $\rho$ from a fixed all- $-1$ vector, concatenates it with dataset embeddings, and maps the decoded representation of $\rho$ to $p(K\mid X)$ .	Uses CIN to map assignment Gram matrices $\{P_{K}^{\top}P_{K}\}_{K=2}^{K_{\max}}$ 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 $\uparrow$	NMI $\uparrow$	$k$ -MAE $\downarrow$	ARI $\uparrow$	NMI $\uparrow$	$k$ -MAE $\downarrow$
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

\Delta_{b}=\mathrm{ARI}_{{\texttt{TabClustPFN}}{}}-\mathrm{ARI}_{b}

. 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

H_{1}:\mathrm{ARI}_{{\texttt{TabClustPFN}}{}}>\mathrm{ARI}_{b}

Figure 14: Wilcoxon signed-rank test results based on per-dataset ARI differences

\Delta_{b}=\mathrm{ARI}_{{\texttt{TabClustPFN}}{}}-\mathrm{ARI}_{b}

. Positive values indicate that TabClustPFN outperforms the corresponding baseline. Colors indicate the one-sided Wilcoxon signed-rank test

p

-value under the alternative hypothesis

H_{1}:\mathrm{ARI}_{{\texttt{TabClustPFN}}{}}>\mathrm{ARI}_{b}

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.

Panel of Figure 17 — 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 $\uparrow$	NMI $\uparrow$	$k$ -MAE $\downarrow$	Time (s) $\downarrow$
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.

Cluster Number Estimation

See all 13 papers

Correcting Class Imbalance in Prior-Data Fitted Networks for Tabular Classification
May 20, 2026Samuel McDowell, Nathan Stromberg, Lalitha SankarClass-Imbalanced Learning Tabular Classification
School of Electrical, Computer and Energy Engineering · Arizona State University, Tempe, AZ
CRUMB: Efficient Prior Fitted Network Inference via Distributionally Matched Context Batching
Jun 9, 2026Jamie Heredge, Mattia J. Villani, Pranav Deshpande +2In-Context Learning Prior-Data Fitted Networks
Global Technology Applied Research, JPMorganChase, New York, NY 10001, USA
TabPFN-MT: A Natively Multitask In-Context Learner for Tabular Data
May 16, 2026Cormac Cureton, Narges ArmanfardTabular Foundation Models Efficient Neural Network Inference
McGill University · Mila - Quebec AI Institute · Montreal, QC, Canada

May 20, 2026J/K move · Enter open · S save

Correcting Class Imbalance in Prior-Data Fitted Networks for Tabular Classification

Samuel McDowell, Nathan Stromberg, Lalitha Sankar

School of Electrical, Computer and Energy Engineering · Arizona State University, Tempe, AZ

Class-Imbalanced Learning Tabular Classification Prior-Data Fitted Networks

Open paper AskPDF arXiv

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.

Abstract

Cluster Number Estimation

Explore similar work

Correcting Class Imbalance in Prior-Data Fitted Networks for Tabular Classification

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

TabPFN-MT: A Natively Multitask In-Context Learner for Tabular Data