EquiTabPFN: A Target-Permutation Equivariant Prior Fitted Network

Comparison with ModernNCA

Thank you for mentioning this method. One reason why we did not include it in the comparison is that the method is not an ICL method and requires being trained on each dataset. However, we agree that it is relevant given that it proposes using a non-parametric output as in our method and we will include it in the paper.

Answers to questions

• Figure 2: We agree with you that the behavior on Figure 2 is surprising for TabPFN-v2. One reason for this behavior is that the model development considered mostly improving point prediction accuracy but has not considered the (non-)equivariant aspect to guide the model development.

• Figure 5: Thank you for raising this point. We apologize for using this Pareto front terminology without properly introducing it, we will clarify in the paper. The Pareto front is defined by the methods that achieve the best tradeoffs of time cost vs performance (%AUC improvement) among the compared methods. These best tradeoff-achieving methods are the ones there is no other method that is strictly better both in terms of performance and time cost.

Implication: The fact that EquiTabPFN + various ensemble sizes, constitute the Pareto front means that: to achieve the highest performance given a fixed time budget, one should use EquiTabPFN + corresponding number of ensembles (closest dot on the Pareto), instead of the other models in the setup we studied.

Thank you for your positive feedback. We provide below answers to your questions, hoping they address them.

Code availability

We unfortunately cannot share the code at the rebuttal as submitting any link is forbidden by the current Neurips policy. We are fully committed to share all the code necessary to reproduce our experiments with the paper including the code of the method and the code to reproduce figures and tables.

Equivariance vs no equivariance

This is a great point that we think should be emphasized in the paper. Here are our current thoughts on the question.

• What does equivariance buy for TabPFN? It allows the model to handle an arbitrary number of classes, which is one of the key strengths of the model. This property would not be true even if the original non-equivariant was trained to achieve 0 equivariance error (for a fixed maximal class number). That being said, one could retain this property (handling arbitrary class number), while breaking equivariance in a controlled way, if it is not always desirable. For example, by adding a form of positional encoding as you suggested (see also answer to Q4). Hence, equivariance, could be a general guiding principle for building models that naturally adapts to varying sizes and shapes. Once constructed, one could choose to break equivariance in a model using general recipes (ex: positional encodings).
• The case of AlphaFold 3: It is unclear from the paper whether the performance boost is due to absence of equivariance, or the use of a better modeling approach (diffusion model instead of an equivariant structure model). The authors explained that they do not use an equivariant diffusion model to simplify the architecture, but they did not claim that using an equivariant diffusion model performed worse. They motivate the use of a diffusion model by noting that the complexity of the structure module in AF2 had only a minor effect on performance and that diffusion models avoid the ‘torsion-based parametrization’ and explicit ‘violation losses on the structure’ present in AF2.

Q1: Computational cost

• Run time comparison: Thank you, we will add those and discuss them in more detail based on the table below showing the average times on the datasets considered in the experiments (either more than 10 classes or less than 10 classes). On datasets with a small number of classes, EquiTabPFN incurs a slowdown of 5x compared to TabPFNv1. This is expected as TabPFNv1 was optimized to handle data with less than 10 classes. On datasets with more than 10 classes, the gap narrows (only 1.3x slowdown) as TabPFNv1 requires ensembling techniques to handle the larger number of classes. A more complete picture accounts for the tradeoff between time cost and performance in figure 5, to which we added TabPFNv1 as requested by reviewer gnTN. It shows that, given a time budget, EquiTabPFN remains the most efficient method most of the time due to the fact that TabPFNv1 typically requires more ensembles to reach good performance likely due to its non-equivariance as illustrated in Figure 3).

• Complexity: The theoretical computational complexity in the number of classes is linear for both, but with a larger factor for EquiTabPFN. The linear dependence for TabPFN comes from projecting the classes into a fixed dimensional vector and then mapping back a fixed dimensional feature to the number of classes, while the dependence on the class number is constant for its backbone. For EquiTabPFN, the backbone features scale linearly in the class number, but it enables handling arbitrary class numbers.

Q2 and Q3: On figure 2 (robustness to class ordering)

From Q2 and Q3, we understand that our current presentation of the figure suggests there could be some symmetries in the decision boundary between classes that the network should capture. We do not expect EquiTabPFN to exhibit such symmetries, as nothing in the model constrains it, e.g. blue vs pink clusters need not have the same shape as nothing would enforce the shape to capture the Voronoi diagram of the 9 training points). We do expect, however, that EquiTabPFN produces the same decision boundaries across the bottom three figures as a consequence of target equivariance (i.e. after permuting the order of the classes numbers). This is indeed the case, up to numerical precision. Thank you for raising this point, we will make sure to clarify this. Q3: Scaling in figure 2: We will clarify that both axes do not have the same scaling.

Q4: Reengineering non-equivariance back into the model

This is a great question! Also related to one from reviewer gntN. Indeed, one could think about breaking equi-variance by providing some form of positional encoding. One would have to adapt the prior to generate non-equivariant but realistic datasets (currently only equivariant datasets are generated). One could then use column index or name, to define the encoding and consider ordinal classes. We believe this can be a promising future work direction.

Q5: Analysis per datasets

Most of the datasets have nominal classes which are equivariant. For the 5/86 datasets that are ordinal, we could not see a particular difference for TabPFN performance. As noted above, we believe it is expected as TabPFN is trained with an equivariant prior so it is unlikely it would be able to exploit non-equivariant information as it was not trained to handle it.

Sorry, I have an additional question. What is the intuition for why this type of equivariance is not used in LLMs? After all, with next-token-prediction, you are performing classification where the number of classes (the vocabulary size) is extremely large, and the ordering of tokens in the vocabulary / embedding table is arbitrary. Assuming that class equivariance has not been attempted for language modeling, why not? Or, if you're aware of previous works where this has been attempted, were there problems scaling this to hundreds of thousands of classes, such that it wasn't worth it? Or perhaps, is the proposed method possibly useful for language modeling?

Thank you for your encouraging and thorough feedback. To address your question on time/performance comparison, we included TabPFN v1 in figure 5 and present the pareto front in table format (below). Below, we address your questions.

Q1 – Tensor shapes

Yes, the tensor shapes are modified as you say. We will add an explanation using shapes in the text for clarity.

Q2 – Reproducing TabPFNv1

Yes, we reused the public reimplementation from the Mothernet paper and we confirm that we could reproduce the results of TabPFNv1 using this setup which resulted in similar performance as the publicly available pre-trained model from the original author paper. We kept the optimization hyperparameters as-is and did not change nor tune them.

Q3 – Averaging TabPFNv1 in figure 2

Yes indeed, TabPFN-V1 converges to better clusters when increasing the number of ensembles. However, as shown in Figure 3, the equivariance error decreases only slowly and the equivariance error is still significant when using 16 ensembles for instance.

Q4 – Critical diagrams (figures 7 & 8)

• Figure 7: Ensembling mechanically helps increase TabPFNv1 performance, but induces extra cost. Please also refer to the updated time/comparison results (answer to Q6), which compares TabPFNv1 with different ensembling to EquiTabPFN.
• Figure 8: Similarly, ensembling increases performance of TabPFNv2 at an extra cost. This is also consistent with the performance vs time cost plot in figure 5 (left).
  Overall these results are consistent with the claim of the paper: EquiTabPFN deals more efficiently (better performance at lower time cost) with an unseen number of classes thanks to its target equivariance architecture.

Q5 – Using TabPFNv2 backbone

We tried TabPFNv2 architecture backbone and modified it to make it target equivariant, then trained it using the publicly available code for the prior used for training TabPFNv1. This led to improvements compared TabPFNv2* (the version we retrained ourselves using publicly available training prior). However, we found that using TabPFNv1’s backbone yielded the best performance overall. We will clarify that in the text.

Q6 – Computational efficiency and memory

• Effect of architecture on context size: Indeed, there is an increase in memory cost, this had an impact on the batch size (first dimension) that we had to reduce during training. However, the context (in terms of the number of samples that can be processed by the model on our devices) was not affected in the experiments. This is likely due to the moderate size of the contexts used whose ranges are within the recommended limits for TabPFN. We agree however, that scaling to larger datasets could cause memory issues. We will clarify this in the text.
• Time/performance comparison with TabPFNv1: Thank you for this comment, it is an excellent point! The focus of the time/performance graph was on TabPFNv2 as the SOTA method. However, as you pointed out, TabPFNv1 is lighter and would benefit from additional ensembling. When adding TabPFNv1 and EquiTabPFN* (the version using TabPFNv2 backbone) to the time/performance graph, we observed that TabPFNv1 indeed contributes (2) points on the Pareto front (2 points), while the points contributed by EquiTabPFN to the front remain unchanged. Hence, EquiTabPFN remains efficient on a wide range of time budgets. As we cannot include the graph here, we present a table summarizing these results. We will add the graph and discussion in the revised version. Time/performance tradeoff: | Pareto points | 1 | 2 | 3 | 4 | 5 | 6 | 7 | |---------|----------|----------|----------|----------|----------|----------|----------| | Old Pareto front | LinearModel | EquiTabPFN(1) |___________| EquiTabPFN(4) | EquiTabPFN(12) | EquiTabPFN(18) | ___________ | | New Pareto front | TabPFNv1(6) | EquiTabPFN(1) | TabPFNv1(18) | EquiTabPFN(4) | EquiTabPFN(12) | EquiTabPFN(18) | EquiTabPFN*(18) | | Time| 0.269 | 0.439 | 0.789 | 1.978 | 3.230 | 4.907 | 9.979 | | Performance| 51.06 | 60.14 | 61.06 | 61.73 | 61.74 | 63.68 | 66.50 |

The table below shows how the pareto front is affected by adding TabPFNv1 and EquiTabPFN* (the version using TabPFNv2 backbone). The first row is the order of the methods crossing the pareto front from left to right. Four observations:

1. Overall, the points forming the pareto and corresponding to EquiTabPFN are unchanged.
2. The leftmost point on the front is no longer LinearModel and is now TabPFNv1 using an ensemble of size 6 (smallest possible to handle more than 10 classes as discussed in the paper).
3. There is an additional point obtained using TabPFNv1 with 18 ensembles, that is inserted between EquitabPFN (1 ensemble) and EquiTabPFN (4 ensembles).
4. An additional point obtained using EquiTabPFN* with 18 ensembles is added at the rightmost side.

General question: Equivariance vs no equivariance

We thank you for this great question, we will discuss this point in the paper. Our thoughts:

• When to (not) be equivariant: We agree that if data is not exchangeable and has additional structure, e.g. ordinal, one should not aim at enforcing equivariance (which would destroy it), and instead leverage the ordinal structure. However, the prior used in TabPFN is equivariant and because of that, the model still ends up learning approximate invariance during training as you pointed out (also shown in Figure 3).
• Controlling (non)-equivariance: It is unclear to us to what extent current models (such as TabPFN) could learn a mechanism to control when to be equivariant given that they are trained with an equivariant prior. However, if equivariance is no longer desirable, one could think of a way for building such a control mechanism in EquiTabPFN to break equivariance while retaining the ability to handle any class number. As pointed out by reviewer Y2PB, a “positional encoding” computed at inference time could be added to the tokens class to differentiate them from each other. When providing the same positional encoding to all tokens, the representation becomes equivariant, otherwise, it is not. For this, the prior would have to be adapted to include examples of non-equivariant and realistic artificial datasets which we believe is an interesting research direction.

Minor points: Estimation by ensembling

Indeed, we will clarify that point. The variance of the estimator could present a dependence on the dimension, however it is challenging for a general function. The closer the function is to equivariance the smaller such variance would be. Hence, the worst scenarios happen when the function is far from being equivariant, in which case, the variance is the highest.

Thank you for your positive feedback. To address your points, we performed the FLOPS and ablations comparison and summarize the findings below.

Flops comparaison. While EquiTabPFN and TabPFN have similar parameter counts, EquiTabPFN uses more FLOPS. On an A100 GPU with 2,000 samples (100 features, 10 classes), EquiTabPFN required 566 GFLOPS vs. 76 GFLOPS for TabPFN (~7.45× more). However, with more classes (15), the gap narrows due to the ensembling needed for TabPFN to handle more than 10 classes (EquiTabPFN: 820 GFLOPS; TabPFN: 456 GFLOPS), consistent with runtime trends.

Still, FLOPS or runtimes need to be traded for performance as in Figure 5 to which we added TabPFNv1 as requested by reviewer gnTN. There, EquiTabPFN offers the best trade-off between time and performance in most cases as summarized in the table below, in particular since TabPFN requires more ensemble to reach good performance. Finally, we note those numbers remain small for a modern GPU given that a single H100 can easily reach 400 TFLOPS on an LLM training workflow for instance.

Table of time/performance tradeoff:

Ablations. Thank you for the suggestion, we will include the following ablation results on different architectures in the revised version. Specifically, we report relative error reduction over TabPFN in two settings: (1) using the TabPFN backbone (without bi-attention) with our equivariant decoder, and (2) using a bi-attention backbone with a standard MLP decoder (as in TabPFNv1). As shown in the table below, the combination of both—biattention and the equivariant decoder, eg the EquiTabPFN model we propose —yields the largest performance gain. The encoder, a simple linear embedding, is considered part of the backbone: fully connected for TabPFN-style models and 1x1 convolution for biattention-based ones.