Multitask-Informed Prior for In-Context Learning on Tabular Data: Application to Steel Property Prediction
Authors: Dimitrios Sinodinos, Bahareh Nikpour, Jack Yi Wei, Sushant Sinha, Xiaoping Ma, Kashif Rehman, Stephen Yue, Narges Armanfard
Organizations: Dep. of Electrical and Computer Engineering, McGill University · Mila – Québec AI Institute · Dep. of Mining and Materials Engineering, McGill University · Algoma Steel Inc.
Accurate prediction of mechanical properties of steel during hot rolling processes, such as Thin Slab Direct Rolling (TSDR), remains challenging due to complex interactions among chemical compositions, processing parameters, and resultant microstructures. Traditional empirical and experimental methodologies, while effective, are often resource-intensive and lack adaptability to varied production conditions. Moreover, most existing approaches do not explicitly leverage the strong correlations among key mechanical properties, missing an opportunity to improve predictive accuracy through multitask learning. To address this, we present a multitask learning framework that injects multitask awareness into the prior of TabPFN--a transformer-based foundation model for in-context learning on tabular data--through novel fine-tuning strategies. Originally designed for single-target regression or classification, we augment TabPFN's prior with two complementary approaches: (i) target averaging, which provides a unified scalar signal compatible with TabPFN's single-target architecture, and (ii) task-specific adapters, which introduce task-specific supervision during fine-tuning. These strategies jointly guide the model toward a multitask-informed prior that captures cross-property relationships among key mechanical metrics. Extensive experiments on an industrial TSDR dataset demonstrate that our multitask adaptations outperform classical machine learning methods and recent state-of-the-art tabular learning models across multiple evaluation metrics. Notably, our approach enhances both predictive accuracy and computational efficiency compared to task-specific fine-tuning, demonstrating that multitask-aware prior adaptation enables foundation models for tabular data to deliver scalable, rapid, and reliable deployment for automated industrial quality control and process optimization in TSDR.
Figures & tables
Fig. 1 : Schematic of Direct Strip Processing Complex
Fig. 2 : An overview of our proposed multitask fine-tuning techniques and the standard TabPFN inference pipeline. (A) Fine-tuning : Xtrain∈RNtrain,D and Ytrain∈RNtrain,T represent the train data and train labels respectively, where Ntrain is the number of training samples, D is the number of features per sample, and T is the number of tasks. To obtain Yˉtrain∈RNtrain,1 , we average (Avg) Ytrain along the task dimension. Using Xtrain and Yˉtrain as inputs to the underlying transformer of TabPFN, we obtain the predicted Yˉ^train . In the standard multitask fine-tuning (M.F.T), we use only the regression loss Lreg as the training signal. In multitask adapter fine-tuning (M.A.F.T), we also use the task-specific losses after obtaining the task-specific predictions Y^train,i∈RNtrain,1 for i∈[1,T] by passing Yˉ^train through an MLP adapter. (B) Inference : We sequentially predict the i th task target by following the standard TabPFN inference pipeline. This involves first updating the prior statistics of our pre-trained TabPFN using the .fit() function with Xtrain and Ytrain,i as inputs, followed by a forward pass using .predict() with Xtest∈RNtest,D as input to obtain Y^test,i∈RNtest,1 , which is the target for the i th task and Ntest is the number of test samples.
Feature Category
Features
Chemical Composition
Al, B, C, Ca, Cr, Cu, Mn, Mo, Nb, Ni, P, S, Si, Sn, Ti, V, N
Caster Section
Avg. Casting Speed, Superheat
Slab Geometry
Slab Width, Slab Gauge, Width (R1 Exit)
Strip Geometry
Final Width
Thickness Measurements
Exit Thickness (R1, F1, F2, F3, F4, F5, F6)
Furnace Parameters
Slab Residence Time
TABLE I : Summary of Features and Targets
Fig. 3 : Spearman correlation matrix of the five mechanical property targets (LYS, OYS, EYS, UTS, ELO). Values represent pairwise rank correlations, with positive correlations shown in red and negative correlations in blue.
Models
Lower Yield Strength
0.2% OS Yield Strength
0.5% Ext Yield Strength
Ultimate Tensile Strength
2” Elongation
ΔM (↑)
MAE% (↓)
PAM 5 (↑)
PAM 2.5 (↑)
EV (↑)
MAE% (↓)
PAM 5 (↑)
PAM 2.5 (↑)
EV (↑)
MAE% (↓)
PAM 5 (↑)
PAM 2.5 (↑)
EV (↑)
MAE% (↓)
PAM 5 (↑)
PAM 2.5 (↑)
EV (↑)
MAE% (↓)
PAM 5 (↑)
PAM 2.5 (↑)
EV (↑)
STL
2.963
83.42
52.69
96.15
2.858
84.96
54.08
96.41
3.180
84.58
51.58
94.85
2.012
94.81
70.16
96.84
3.762
73.85
43.30
90.00
+0.00
MTL
3.026
83.09
50.16
96.11
2.921
82.25
52.27
96.40
3.187
83.95
50.74
94.89
2.099
93.79
67.88
97.05
3.987
69.81
41.29
88.89
-2.91
MMoE
2.933
84.11
53.60
96.03
2.836
85.29
54.53
96.36
3.148
84.08
52.81
94.89
2.018
94.32
70.46
97.15
3.790
72.43
43.36
89.81
+0.35
PLE
2.924
84.26
53.82
96.06
2.819
85.03
55.33
96.43
3.164
83.52
52.96
94.57
1.980
94.59
71.20
97.15
3.826
72.25
43.03
89.57
+0.61
STEM
2.927
84.16
52.58
96.27
2.848
84.89
54.07
96.21
3.126
83.87
52.19
94.67
2.014
94.19
69.95
97.15
3.956
70.77
41.47
88.80
-0.40
TABLE II : Performance Comparison Across Models and Tasks
Fig. 4 : Multitask gain ( Δm ) as a function of fine-tuning time budget for four strategies: no fine-tuning (n.f.t), single-task fine-tuning (s.f.t), multitask fine-tuning (m.f.t), and multitask adapter fine-tuning (m.a.f.t).
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
Recent progress in tabular foundation models suggests that training on synthetic tasks can substantially improve in-context learning capabilities, with overall performance largely depending on how well models can infer task-specific predictive relationships from the available context during inference. In this paper, we introduce TAFFY, a tabular foundation model with an In-Context Diversity Prior and a Task-Conditioned Looped Transformer that strengthen this ability. Specifically, to construct each synthetic pretraining context, the In-Context Diversity Prior samples from multiple related environments derived via controlled interventions and distribution shifts on a shared causal process. This in-context diversity encourages the model to learn a more comprehensive and task-specific representation. Moreover, the Task-Conditioned Looped Transformer iteratively and selectively applies a shared group of Transformer blocks to refine contextual representations, with a task-conditioned gate modulating the final hidden-state update. This enables task-adaptive iterative refinement. Together, these components encourage the model to identify predictive relationships from contextual contrasts during pretraining and dynamically modulate context integration for each task. Across six classification and five regression benchmark datasets, TAFFY attains the lowest average rank.
Zijian Li, Xiangchen Song, Gongxu Luo +7
Carnegie Mellon University · Mohamed bin Zayed University of Artificial Intelligence · Guangdong University of Technology +1
Tabular foundation models enable in-context learning by predicting unseen observations from a labeled support table without task-specific parameter updates. Despite recent advances, most existing approaches rely on dense architectures that apply the same transformations across features and datasets. We introduce OrionMoE-Small-v1.0, a compact tabular foundation model that combines distribution-aware representations with dataset-conditioned sparse mixture-of-experts computation. Built on a dual-axis transformer architecture, OrionMoE-Small selectively activates experts with routing decisions informed by both individual cell representations and the broader dataset context. This design enables adaptive allocation of computation across heterogeneous tabular problems while retaining a shared computational pathway across tasks. The model is pretrained entirely on synthetic classification tasks generated from a structured causal prior and performs downstream prediction without gradient-based adaptation. Experiments across multiple tabular benchmarks demonstrate competitive performance against substantially larger foundation models, with particularly strong results on small-context and high-dimensional datasets. Further analysis reveals favorable relative performance under class imbalance and highlights the importance of feature composition.