Multitask-Informed Prior for In-Context Learning on Tabular Data: Application to Steel Property Prediction

Dimitrios Sinodinos; Bahareh Nikpour; Jack Yi Wei; Sushant Sinha; Xiaoping Ma; Kashif Rehman; Stephen Yue; Narges Armanfard

cs.LGMar 24, 2026

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.

Abstract

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 : $X_{\text{train}}\in\mathbb{R}^{N_{\text{train}},D}$ and $Y_{\text{train}}\in\mathbb{R}^{N_{\text{train}},T}$ represent the train data and train labels respectively, where $N_{\text{train}}$ is the number of training samples, $D$ is the number of features per sample, and $T$ is the number of tasks. To obtain $\bar{Y}_{\text{train}}\in\mathbb{R}^{N_{\text{train}},1}$ , we average (Avg) $Y_{\text{train}}$ along the task dimension. Using $X_{\text{train}}$ and $\bar{Y}_{\text{train}}$ as inputs to the underlying transformer of TabPFN, we obtain the predicted $\hat{\bar{Y}}_{\text{train}}$ . In the standard multitask fine-tuning (M.F.T), we use only the regression loss $\mathcal{L}_{\text{reg}}$ 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 $\hat{Y}_{\text{train},i}\in\mathbb{R}^{N_{\text{train}},1}$ for $i\in[1,T]$ by passing $\hat{\bar{Y}}_{\text{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 $X_{\text{train}}$ and $Y_{\text{train},i}$ as inputs, followed by a forward pass using .predict() with $X_{\text{test}}\in\mathbb{R}^{N_{\text{test}},D}$ as input to obtain $\hat{Y}_{\text{test},i}\in\mathbb{R}^{N_{\text{test}},1}$ , which is the target for the $i$ th task and $N_{\text{test}}$ 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				$\Delta_{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 ( $\Delta_{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).

Multi-Task Learning

Multi-task learning trains models to solve related tasks together by sharing representations, parameters, or supervision.

See all 389 papers

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
TAFFY: A Task-Adaptive Tabular Foundation Model with In-Context Diversity
Oct 6, 2026Zijian Li, Xiangchen Song, Gongxu Luo +7Synthetic Data Pretraining Tabular Foundation Models
Carnegie Mellon University · Mohamed bin Zayed University of Artificial Intelligence · Guangdong University of Technology +1
OrionMoE-v1.0: A Mixture-of-Experts Foundation Model for Tabular In-Context Learning
Oct 14, 2026Mohamed Bouadi, Nassim Bouarour, Aditya Tanna +1Tabular Foundation Models Tabular ICL
Lexsi Labs

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

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

Cormac Cureton, Narges Armanfard

McGill University · Mila - Quebec AI Institute · Montreal, QC, Canada

Tabular Foundation Models Efficient Neural Network Inference Multi-Task Learning Tabular ICL

Open paper AskPDF arXiv

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)

O(1)

forward passes, offering a massive computational efficiency improvement for multi-target tabular applications.

Abstract

Multi-Task Learning

Explore similar work

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

TAFFY: A Task-Adaptive Tabular Foundation Model with In-Context Diversity

OrionMoE-v1.0: A Mixture-of-Experts Foundation Model for Tabular In-Context Learning