Title: Drift-Resilient TabPFN: In-Context Learning Temporal Distribution Shifts on Tabular Data
Venue: NeurIPS 2024
Accepted: yes
Keywords: Temporal Distribution Shifts

In-Context Learning

Bayesian Inference

Prior-Data Fitted Networks

Temporal Domain Generalization

Structural Causal Models

TabPFN

Tabular Data Modeling

Out-Of-Distribution Generalization

Domain Generalization

Meta-Learning

Concept Drift
Abstract: While most ML models expect independent and identically distributed data, this assumption is often violated in real-world scenarios due to distribution shifts, resulting in the degradation of machine learning model performance. Until now, no tabular method has consistently outperformed classical supervised learning, which ignores these shifts. To address temporal distribution shifts, we present Drift-Resilient TabPFN, a fresh approach based on In-Context Learning with a Prior-Data Fitted Network that learns the learning algorithm itself: it accepts the entire training dataset as input and makes predictions on the test set in a single forward pass. Specifically, it learns to approximate Bayesian inference on synthetic datasets drawn from a prior that specifies the model's inductive bias. This prior is based on structural causal models (SCM), which gradually shift over time. To model shifts of these causal models, we use a secondary SCM, that specifies changes in the primary model parameters. The resulting Drift-Resilient TabPFN can be applied to unseen data, runs in seconds on small to moderately sized datasets and needs no hyperparameter tuning. Comprehensive evaluations across 18 synthetic and real-world datasets demonstrate large performance improvements over a wide range of baselines, such as XGB, CatBoost, TabPFN, and applicable methods featured in the Wild-Time benchmark. Compared to the strongest baselines, it improves accuracy from 0.688 to 0.744 and ROC AUC from 0.786 to 0.832 while maintaining stronger calibration. This approach could serve as significant groundwork for further research on out-of-distribution prediction.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper aims to address the temporal distribution shifts in Tabular data based on TabPFN. Concretely, the proposed method use two structural causal models (SCM) to model the prior for TabPFN, one for the gradual shift of inductive bias over time and the other is to model the shift of the first SCM. Empirical studies on various datasets illustrate the effectiveness of the proposed approach.\n\nSTRENGTHS:\n1. The paper aims to address temporal shifts in tabular data, which could be a useful setting in practice.\n2. The proposed method intuitively makes sense.\n3. Experiments on synthetic and real datasets have been tested. The code is released.\n\nWEAKNESSES:\n1. As far as I am concerned, CatBoost, XGboost, and LightGBM do not explicitly consider the temporal shift. The proposed method introduced additional sampled data from a temporal shift prior to PFN. I would doubt such a comparison is not that fair. Does any of the baselines in [1]  outperform the proposed method on the datasets used in this paper?\n2. The experimental setting could be more detailed in the main paper.\n3. It seems the proposed method is a simple modification to TabPFN that changes the SGM prior with additional temporal shift add-ons. So it's a little bit combinatorial to me, but not a huge problem.\n4. Modelling the temporal shifts of inductive bias shares similarity with some continual learning approaches using meta-learning techniques, which lack a brief discussion, such as [2, 3, 4].\n\n[1] Benchmarking distribution shift in tabular data with tableshift. Advances in Neural Information Processing Systems,36 (2024).\n\n[2] Online fast adaptation and knowledge accumulation (osaka): a new approach to continual learning. Advances in Neural Information Processing Systems, 33:16532–16545, 2020.\n\n[3] Reconciling meta-learning and continual learning with online mixtures of tasks. Advances in Neural Information Processing Systems, 32, 2019.\n\n[4] On the stability-plasticity dilemma in continual meta-learning: theory and algorithm. Advances in Neural Information Processing Systems, 36 (2024).\n\nQUESTIONS:\n1. Why do you need the functional representation graph $\\tilde{\\mathcal{G}}$? What are the benefits of the original sampled graph?\n2. Why do sparse shifts allow for causal reasoning (Line 195)? \n3. How are the edge weights decided?\n4. Why can SCM with NN-based causal mechanisms and nonlinear activations extrapolate values out of data distribution (Line 208)? Is this due to a causal mechanism or NN?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose an extension of the TabPFN framework such that the model can be resilient to distribution shifts during the inference phase by observing the shifted samples in-context. Akin to TapPFN, the pipeline involves pretraining based on an SCM, and the authors introduce a temporal aspect to the SCM construction such that the model will be aware that several common types of distribution shifts can be possible at inference time.\n\nSTRENGTHS:\nThe paper was well structured and well written. The motivation of the work was described in a compelling way. The objective of the work was clearly described. The description of the SCM was sufficiently clear and intuitively it makes sense that the described approach should lead to an improvement in handling cases where a distribution shift is noted. All the figures were helpful in understanding the details of the method, and I found the visualization of the decision boundaries especially convincing. It seems there could be a real world impact in areas where TabPFN-like methods are a good fit based on this work.\n\nWEAKNESSES:\nI did not feel the work had any major structural weaknesses within the scope of what the authors wished to achieve.\n\nA small suggestion that the title should imply that the work is about temporal domain shifts.\n\nThe “runs in seconds” clearly would not be true for all kinds of target scales, although it rightfully highlights the potential of the approach in some scenarios. I recommend adding a short qualifier on the scale where this is true.\n\nAdditional clarity on what exactly is the proposed input and output to the final model would have been appreciated in the methods section. For instance, does the model need to observe all previous intervals eg. [0 1 2] and the future intervals [3 4 5]? Only the previous intervals? Or just a trailing list of recent intervals? What is the justification that future intervals [4 5] will be available to the user if they are still in the present (domain 3)? \n\nDue to some of the ambiguity, it was also not clear at a high level if a user needs to directly indicate to the model that a domain shift is generally expected at certain domain indices, or the model makes such a decision for every domain index internally.\n\n“We approximated domain indices … based on features that encode temporal information, which we transformed into discrete intervals.” It seems that the accuracy attained by the models could vary a bit based on the way that the domains are approximated, unless the discretization was completely random uniform and guaranteed to be fair. Some additional assurance seems deserved on how this approximation was done.\n\nI think the authors motivate the problem enough such that readers will be convinced that existing methods would not do well under distribution shifts. Given this, I think a more competitive benchmark could be to see how well baselines trained on small parts of the OOD do well on the OOD domains, and how this compares with the various ways of applying DSTabPFN. I would be curious to know if this was considered.\n\nQUESTIONS:\nDo the authors have any more considerations on the extent to which in-domain (ID) capabilities are affected when adopting the Distribution Shift TabPFN approach as opposed to the original TabPFN? This is hinted in line 287 but it seemed like an important point that gets at some quantitative and motivation-related questions. For instance, what is the false positive rate of the DSTabPFN guessing a distribution shift when there is none? What if there are some outliers in the inference phase despite the underlying distribution will continue longer term? Should users rely on the model to distinguish between one shift to another or only turn to DSTabPFN when a distribution shift is noticed? Some more intuition on this matter may convince readers further.\n\nWhat is the rationale for the compute budget imposed in line 262? Generally eg. XGB would be parallelized with dozens of cores and requires a long time to cross validate & fit, but perhaps this was not necessary for the scale of the datasets explored.\n\nIt would be great to understand how the approach scales with the amount of samples from the SCM. For instance, are x5 more samples (taking x5 times longer to train) needed for the model to be robust to 5 different types of distribution shifts? Or is a similar accuracy to TabPFN possible with an equal number of samples while also being robust to shifts? Is the reason why DSTabPFN does not match the accuracy of TabPFN because it crosses into an overfit regime with too many samples? \n\nA related analysis that could have been nice to have is to see how many “shots” in the shifted distribution are needed for DSTabPFN to adjust appropriately.\n\nLIMITATIONS:\nThe work discusses limitations to a fair extent in the final section. I did not assess that the work could have a major negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a modification of TabPFN that incorporates a novel prior to incorporate temporal shifts. In particular, the proposed method introduces an additional SCM that, through a temporal representation (Time2Vec), learns to modify the model parameters in response to shifts. The authors train a version of the proposed model and compare it to TabPFN, GBDT baselines, and standard ERM and SWA baselines.\n\nOverall, the paper's presentation is good (but could be improved), and this appears to be a logical addition to the TabPFN method. A method capable of zero-shot transfer to new distributions seems especially well-motivated, given that no labeled data from the target domain is available in many real-world scenarios. However, some revisions and additions to the paper are badly needed. For example, the literature review is inadequate and needs a complete rewrite; there are no domain generalization or robustness methods included in the experimental results (besides the proposed method), and the ablation study is too limited (in fact, the current results suggest that there is no advantage to the Time2Vec addition, which raises questions about the overall approach). I think that with extensiv revisions this paper could be brought up to acceptance level and am open to raising my score, but the scope of additions and revisions feels on the borderline of too much for a simple camera-ready phase.\n\nSTRENGTHS:\nFull review here:\n\n# Major comments\n\n* The paper is framed in a way that I find somewhat misleading: the title, abstract, and main text repeatedly frame the current work as an approach to \"distribution shift\". However, the approach is limited to strictly *temporal* shifts, with no discussion or demonstration of its extensibility to non-temporal shifts. This introduces some unnecessary confusion into the paper and its potential applications -- please either clarify this in the title/abstract/text (preferred), or add a thorough discussion why this should be considered a general method for distribution shift (less preffered unless the authors can already demonstrate this capacity, in which case the current framing does indeed fit).\n\n* I like that the paper unifies all shifts (covariate, label, concept) under a single framework. This is a limitation of some (but not all) existing approaches to domain generalization, and it seems to be a clear advantage of this approach which the authors rightly highlight. However, it would also be useful to see a set of controlled experiments which separately vary these different components, to understand how the model's performance varies as different forms of shift are present; these impacts likely depend on the model's pretraining distribution in subtle ways).\n\n* The \"related work\" section neither surveys relevant related work at a level even close to comprehensive, nor does it actually discuss related works relevant to understanding the current paper. In my opinion, this section should be completely rewritten. For example:  (1) the paper does not discuss related benchmarking or empirical studies such as Shifts [1] and Shift 2.0, Wild-Tab [2], TableShift [3] (already cited but not mentioned in related work despite applicable findings), and WhyShift [4]. All of these should be discussed and related to the current work. (Concurrent work TabRed [8] seems also relevant, but of course was published after this submission.) (2) the paper does not mention at all the many previously-proposed methods for robust learning and domain generalization, including those covered in the previous work. The only methods mentioned, strangely, are two methods that the authors state are *excluded* from the current study. (3) Many works have recently proposed changes to or conducted  empirical studies of TabPFN e,g, [5, 6, 7]; it would be useful to position the current work in relation to these studies.\n\n* *Ablation study*: The ablation study should be expanded, and discussed in the main text. In particular: (1) the current results show no benefit from the Time2Vec addition, which seems to suggest that it shouldn't be included in the model; (2) the authors seem to have only conducted a single run, of the ablated model despite this evidence -- a number of trials equal to the other experiment seems in order. (3) please discuss the results in the main text, even if only in a sentence or two that state the conclusion from the experiment.\n\n* *Resources for baselines*: The authors restrict all baselines to 1200 seconds on 8CPU/1GPU, but the proposed method is trained for 30 epochs on 8 GPUs with a random search over 300 configurations. This setup seems to starve the baselines of resources. Please either provide evidence that the baselines have in fact saturated (i.e. they would not be improved from further tuning) or allocate similar resources to the baselines as to the main model -- otherwise it is impossible to distinguish whether the proposed method simply benefits from more compute + hyperparameter tuning vs. whether it is in fact an improvement.\n\n* *Lack of relevant baselines:* I was quite surprised to see that the authors do not include any domain generalization or robustness methods in their study. Open-source implementations of many relevant methods are available, including (I believe) through Wild-Time and TableShift. Please comment on why these were not included (less preferred) or add them to the experimental results (preferred), also conducting fair hyperparameter tuning as described above.\n\n# Minor comments\n\n* More details on the data should be in the main text. What are the datasets, where are they drawn from, why do the authors use 8 synthetic and 10 real-world? None of these are given, and so it is hard to contextualize the results or assess the reliability of the empirical study.\n\n* It is interesting that CatBoost performs quite strongly on the proposed methods, particularly ID on the synthetic data. This could perhaps be worth mentioning in the main text, as it is in line with the findings of other tabular studies (i.e. CatBoost tends to slightly outperform XGBoost and LightGBM).\n\n# Typos etc\n\n* Abstract: \"with even smaller gains for tabular data\" - this is unclear, please revise.\n\n* L48-49: \"Building on...\" this sentence is a fragment, please revise.\n\n* Table 1: I don't see the acronym \"SWA\" defined in the paper (ERM is not defined either, but this is more familiar to most readers).\n\n* L303: there are 2 lines hanging onto Page 9; the figure/table positioning should be changed to avoid this.\n\n# References\n\n[1] Andrey Malinin, Neil Band, Yarin Gal, Mark Gales, Alexander Ganshin, German Chesnokov, Alexey Noskov, Andrey Ploskonosov, Liudmila Prokhorenkova, Ivan Provilkov, et al. Shifts: A dataset of real distributional shift across multiple large-scale tasks. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2022.\n\n[2] Kolesnikov, Sergey. \"Wild-Tab: A Benchmark For Out-Of-Distribution Generalization In Tabular Regression.\" arXiv preprint arXiv:2312.01792 (2023).\n\n[3] Gardner, Josh, Zoran Popovic, and Ludwig Schmidt. \"Benchmarking distribution shift in tabular data with tableshift.\" Advances in Neural Information Processing Systems 36 (2024).\n\n[4] Liu, Jiashuo, et al. \"On the need for a language describing distribution shifts: Illustrations on tabular datasets.\" Advances in Neural Information Processing Systems 36 (2024).\n\n[5] Breejen, Felix den, et al. \"Why In-Context Learning Transformers are Tabular Data Classifiers.\" arXiv preprint arXiv:2405.13396 (2024).\n\n[6] Feuer, Benjamin, et al. \"TuneTables: Context Optimization for Scalable Prior-Data Fitted Networks.\" arXiv preprint arXiv:2402.11137 (2024).\n\n[7] Ma, Junwei, et al. \"In-Context Data Distillation with TabPFN.\" arXiv preprint arXiv:2402.06971 (2024).\n\n[8] Rubachev, Ivan, et al. \"TabReD: A Benchmark of Tabular Machine Learning in-the-Wild.\" arXiv preprint arXiv:2406.19380 (2024).\n\nWEAKNESSES:\nSee above.\n\nQUESTIONS:\nSee above.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the setting of non-iid train / test distribution shift in the tabular machine learning. The authors extend the previous TabPFN sota tabular model to deal with domain shift cases. TabPFN uses SCM to model data prior, and in this paper, the core idea is to update the SCM prior graph by hypernetworks. The output of hypernetworks is the SCM update. As a result, TabPFN could deal with data shift. The whole idea is very simple and the authors demonstrate the superior effectiveness in various synthetic and real datasets\n\nSTRENGTHS:\n- The authors study the data distribution shift in the tabular domain, which is not commonly studied and provides practical value\n- The authors present a simple method (hypernetwork) to update SCM prior \n- The resulted method has been verified in numerical experiments with synthetic and real datasets\n\nWEAKNESSES:\nThe main weakness lies in the experimentation. See the questions below. Technically, the core part is the integration of hypernetworks, which seems very simple. However, in the experiments, there is little discussion of this module. Besides, since many synthetic datasets are used, it is nature to study the SCM from a more explicit way and especially how hypernetworks work well in the setting.\n\nQUESTIONS:\n- If I understand correctly, the core part of the method lies in training a hypernetworks to update SCM. All the rest follows TabPFN?  \n- It is vague how hypernetworks is trained and processed. It is suggested to provide details on this\n- The authors mentioned in Figure 4, various cases of distribution shifts, how are these shifts mitigated in the experiments? Any investigation on different cases?\n- How good has the hypernetworks been trained? Since there is many synthetic datasets used, it is great to study cases where hypernetworks updated the SCM correctly\n- How much data is required to train hypernetworks? \n- How is the hypernetworks training related to number of features, number of classes, etc.  since TabPFN has limitations on the dataset requirement.\n- In Table 1, it is recommended to separate real datasets and synthetic datasets. It is also necessary to briefly discuss the nature of the used datasets\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}