Title: EquiTabPFN: A Target-Permutation Equivariant Prior Fitted Network
Venue: NeurIPS 2025
Accepted: yes
Keywords: tabular data

equivariance

foundation model

TabPFN
Abstract: Recent foundational models for tabular data, such as TabPFN, excel at adapting to new tasks via in-context learning but remain constrained to a fixed, pre-defined number of target dimensions—often necessitating costly ensembling strategies. 
We trace this constraint to a deeper architectural shortcoming: these models lack target-equivariance, so that permuting target-dimension orderings alters their predictions. This deficiency gives rise to an irreducible “equivariance gap,” an error term that introduces instability in predictions. We eliminate this gap by designing a fully target-equivariant architecture—ensuring permutation invariance via equivariant encoders, decoders, and a bi-attention mechanism. Empirical evaluation on standard classification benchmarks shows that, on datasets with more classes than those seen during pre-training, our model matches or surpasses existing methods while incurring lower computational overhead.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nTabPFN is inherently limited to predicting no more than 10 classes, as Hollmann et al. (2025) observe:\n\n> “Thus, natively our method is limited to predicting at most 10 classes.”\n\nThe authors of the current paper argue that this ceiling arises from the model’s lack of equivariance with respect to the target dimension. To overcome it, they introduce an architecture that projects each target dimension into its own token, thereby avoiding the identified weakness. On a large suite of datasets, they empirically demonstrate the resulting performance gains.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n\n1. **Significance.** The authors propose a relatively straightforward modification that performs exceptionally well in practice.\n2. **Theoretical justification.** By introducing the notion of an *equivariance gap*, the authors show that methods lacking the equivariance property are theoretically bound to underperform.\n3. **Extensive empirical validation.** The paper presents a comprehensive evaluation: both in simulation and on real-world dataset, highlighting key design choices, pointing out potential flaws, and reporting performance alongside computational cost for all methods considered.\n4. **Clarity.** The presentation is clear and well-structured; multiple plots and tables effectively communicate the core ideas.\n\n## Weaknesses\n\n1. **Originality.** The final prediction layer, $\\sum_{n=1}^N y_n \\omega_n$ (without the residual MLP, where weights $\\omega_n$ are computed via attention), together with the comparison to KNN, suggests a close relationship to nearest-neighbor methods. It would be valuable to include a comparison with ModernNCA (Ye et al., 2025), which employs a similar prediction layer, even though it is not a foundation model and must be trained per dataset.\n\n---\n\nYe, Han-Jia; Huai-Hong Yin; and De-Chuan Zhan. “Modern Neighborhood Components Analysis: A Deep Tabular Baseline Two Decades Later.” *ICLR* (2025).\n\nQUESTIONS:\n1. **Figure 2.** I was surprised to see TabPFN-v2 performing worse than the original TabPFN, even though your architecture appears conceptually closer to TabPFN-v2 and achieves results closer to or better than TabPFN. Could you clarify why TabPFN-v2 lags behind TabPFN?\n\n2. **Figure 5.** You refer to a “Pareto front,” but the concept is not explained in the text. Could you elaborate on what a Pareto front represents in this context and any consequences stemming from it?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nEquiTabPFN is a proposed modification to TabPFN, a transformer model that has been meta-learned to perform in-context learning based tabular classification. In particular, it fixes the lack of output class order equivariance in TabPFN, with evidence from both theoretical analysis and from tabular benchmarks that this produces meaningful improvements.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\n(S1) [Significance, Originality] The presented problem (lack of equivariance for permutions of classes in TabPFN) and proposed solution (a classification head that is class order equivariant) are relevant, clever, and interesting. \n\n(S2) [Significance, Quality, Originality] The theoretical analysis showing the benefit of equivariance is very nice. Instead of being a box-checking exercise or an impress-the-reviewers exercise, the analysis proves something interesting and contributes to the narrative.\n\n(S3) [Significance, Quality] The results on both the initial analyses and also on the benchmarks are great. It's surprising that ensembling in TabPFN fails so badly to attain (approximate) equivariance, and that the very rough decision boundaries in Figure 1 would be solved by equivariance.\n\n(S4) [Clarity] The paper is very clearly written.\n\nWeaknesses\n\n(W1) [Significance] Code was not provided, but the authors explained that anonymization wasn't ready at time of submission. So we hope that this can be addressed by the rebuttal period.\n\n(W2) [What follows is a subjective, debatable point, and is not a reason for rejection.] Efforts to introduce invariances / equivariances have generally faced the Bitter Lesson. For example, the most recent version of AlphaFold found that dropping rotational equivariance offered better performance. Why should we not suspect that EquiTabPFN won't meet the same fate? For example, one might imagine that in many cases, modelers sort their classes (with say the first one being the \"default\" majority class). As another example, certain classification problemsPerhaps a non-equivariant model will be able to learn and exploit such phenomena while defaulting to equivariance.\n\nQUESTIONS:\nQuestions:\n\n(Q1) I see that EquiTabPFN was trained with similar number of parameters as TabPFN for correctness evaluation fairness, so I am guessing that they are probably similar in computation time as well. But could you please comment more on this (and ideally include in the paper)? Did EquiTabPFN produce any slowdown in runtime? Are the computational complexities the same in terms of number of classes?\n\n(Q2) Even though EquiTabPFN is equivariant by construction, it seems to me from Figure 1 that it's not perfectly equivariant. For example the pink and blue classes aren't identical in shape, with the pink boundary curved and the blue boundary pointy. Is this due to numerical imprecision?\n\n(Q3) Are the x-axis and y-axis in Figure 1 equally scaled? It's not clear to me what are the sets of classes that should be identical in decision boundary shape to each other.\n\n(Q4) Related to (W1), is there a way to reengineer non-equivariance back into the model, the same way position encodings allow one to \"learn to bypass\" attention's permutation-invariance?\n\n(Q5) Analysis about which datasets in the benchmark improved or worsened from TabPFN to EquiTabPFN would be interesting. Were there any datasets that have some kind of class-order non-equivariance that TabPFN manages to exploit?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a mechanism to make TabPFN-style models equivariant to the ordering of the class labels, something which is often achieved in practiced by randomizing over permutations of those.\nThe authors methods is to essentially tie a vector to each class dimension and perform a cross attention at the last layers to get similarity between the \"validation\" token refined through the forward pass and the class tokens. This also allows them to extend naturally beyond a fixed number of classes which is desirable.\n\nSTRENGTHS AND WEAKNESSES:\nSimple architectural improvements that can improve the models are always welcome, I have a few questions and reservations.\n\nQ1: To make sure I understand correctly, your activations vector are of shape $(B, N+M, q+1, d)$? For the instance attention you reshape to $(B*(q+1), N+M, d)$ where batch size is first, then transformer sequence length, then transformer dimension. For the self attention across labels dims/features, it would be $(B*(N+M), q+1, d)$? \n\nThe method overall makes sense, and so does the cross-attention to compute outputs.\nOn the evaluation side, I appreciate the ablations/baselines with variants of TaPFN v1/2 to disentangle the impact of different choices.\n\nQ2: I see that you have tried to be faithful to TabPFN v1 procedure and hyperparameters, so would it be fair to say that using your training code but using TabPFN v1 architecture you would get very similar experimental results to TabPFN v1? I am asking because some small choices (adding a cosine in the prior bank of activations, changing the some optimization parameters...) can affect/improve the performance of those model, and as such I am not totally certain I can use TabPFN v1's performance as a reference point for equitabpfn.\n\nQ3: I am not too sure how a 9 instances predictions problem is really all that insightful but the vizualization is nice. However as TabPFN v1 is typically averaged over class permutation, does this averaging actually lead to decent clusters? \n\nQ4: could you comment on Figure 7 and 8 in appendix? it seems that for <=10 classes equitabpfn has a performance between tabpfn v1 and tabpfn v1 ensemble?\n\nQ5: Apologies if I missed but have you tried your method on top of tabpfn v2? What are the results?\n\nQ6: *Computational efficiency and memory* Finally this is my most important remark and question. If I understand correctly, compared to tabpfn v1, you now deal with a 4D tensor, and as such your activations are much larger (q+1 times) which is bound to limit the context size you can use. Does your method exhibit a reduced context size compared to v1?\nFurthermore, it does also increase the computations if I understand correctly. Where is TabPFN v1 and TabPFN v1-ensemble in the time/performance graph? It seems to me that for the same memory and computational cost you can afford to have the vanilla more efficient architecture+some bagging for free.\n\nLastly, this is a more general open question on which I am interested to have your opinion. I think being able to extend beyond a fixed number of classes is great. I am not convinced baking equivariance in the model is necessary. In the same way that Alphafold 2 simply used data augmentation over many protein orientations instead of baking the equivariance in the model, a sufficiently large model is able to learn invariance/equivariance on its own. Even more that this, there may be some patterns in the training data that could be destroyed by equivariant/invariant transformations. In tabular data, classes are actually often ordinals in disguise (age, number of children etc...) so I tend to believe that just showing real data to a sufficiently powerful model would enable it not only to be invariant to order but also *when* to be.\n\n## Minor points\nLine 276: While indeed O(q!) would be needed for perfect invariance to ordering, in practice the prediction would concentrate with O(1/sqrt(N_ensembles)) which should converge on a given class in actually a not so high number of ensembles.\n\nQUESTIONS:\nSee above\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

comments: SUMMARY:
This paper introduces an equivariant model for tabular data to address the lack of target equivariance in existing methods. It does this by adding equivariance to the components of a TabPFN model, especially by adding a bi-attention mechanism in the backbone of the architecture. It demonstrates the superiority of the model over existing baselines in the case of unseen class counts at test time.

STRENGTHS AND WEAKNESSES:
Strengths:
- The paper addresses an important challenge with an original methodology.
- It presents a theoretical argument for their approach, and shows that target equivariant models are optimal. It shows that models not inherently designed to be equivariant like so are either less robust or use more computational power to achieve equivariance. 
- It empirically shows that the proposed model is more robust to the number of target classes, performing better than current models in the case that test classes differ from pre-training. 



Weaknesses:

The paper already addresses the limitation that the attention mechanism adds quadratic complexity. In the interest of fairness, they compare to baselines with similar number of model parameters. But in keeping with LLM literature, it would be a fairer comparison to compare compute flops to baselines, especially given that their method adds more complexity not necessarily reflected in the amount of parameters.

QUESTIONS:
- What is the difference in flops between TabPFN and EquiTabPFN?
- What is the benefit of the individual equivariance adding components in the encoder, backbone, and decoder? An ablation would be interesting to see here. It is possible that one component is more important than others and can help reduce the complexity of the method.

LIMITATIONS:
Yes
TITLE: Official Review
CLARITY: 4
IMPACT: 4
ORIGINALITY: 4
DATE: 2025-07-02