MatrixFormer: A Foundation Model for Matrix Completion

Dwaipayan Saha; Jacob Feitelberg; Kyuseong Choi; Raaz Dwivedi; Anish Agarwal

cs.LGOct 5, 2026

MatrixFormer: A Foundation Model for Matrix Completion

Authors: Dwaipayan Saha, Jacob Feitelberg, Kyuseong Choi, Raaz Dwivedi, Anish Agarwal

Organizations: Cornell Tech

Abstract

Matrix completion underlies problems from tabular imputation to causal inference, yet existing tabular foundation models treat it as entry-by-entry prediction, repeating context for every target and discarding the matrix's two-dimensional structure. We introduce MatrixFormer, a pre-trained matrix-native transformer that predicts a full distribution for every missing entry in a single forward pass. MatrixFormer is trained entirely on synthetic low-rank and latent-factor matrices under diverse missingness patterns. Applied zero-shot and with the same model weights, MatrixFormer achieves competitive performance on causal inference panel-data tasks, language-model benchmark-score completion, tabular imputation, and recommendation systems matrix completion. These results position MatrixFormer as a general-purpose foundation model for matrix completion.

Figures & tables

Figure 1: MatrixFormer forward pass. Each matrix cell is represented by a token with an explicit missingness embedding. Missing values use a learned mask token. Learned row- and column-CLS tokens are added around the matrix, and each transformer block alternates feature attention across columns within a row and datapoint attention across rows within a column. After

L

blocks, the CLS border is removed and a decoder predicts the masked entries.

Figure 2: Fixed-column inference runtime for legacy TabImpute and MatrixFormer, with logarithmic row and runtime axes.

Figure 3: Normalized causal-panel scores (lower is better; 0 = best and 1 = worst of the seven methods shown in a draw; dashed line at 1). The blue bars show the MatrixFormer ensemble. Top: RMSE on the hidden treated block. Bottom: absolute error of the average treatment effect on the treated. Bars are means over draws with designs weighted equally; whiskers are 95% intervals over draws.

Figure 4: MissBench: Overall, MCAR, MAR, and MNAR. Columnwise NRMSE across 42 datasets at the nominal 40% setting; lower is better. MatrixFormer is blue; baselines are gray. Methods match the UCI figure, with TabImpute (legacy) additionally shown here. The overall panel includes all eleven mechanisms. Rows are sorted by overall NRMSE, with that order shared by the four panels. Whiskers are 95% dataset-bootstrap confidence intervals.

\dagger

denotes incomplete coverage.

Figure 5: UCI: Overall, MCAR, MAR, and MNAR. The suite contains twelve datasets, five missingness rates, and ten repeats. Each mechanism averages available rates and repeats within datasets, then weights the available datasets equally; lower is better. The overall panel includes all three mechanisms. Whiskers are 95% dataset-bootstrap confidence intervals. Rows follow UCI overall NRMSE.

\dagger

denotes incomplete coverage.

	MovieLens 100K		Netflix small
Method	RMSE	MAE	RMSE	MAE
MatrixFormer	0.8874	0.7279	0.8617	0.7330
SVD	0.8998	0.7430	0.8621	0.7358
PMF	0.9138	0.7534	0.8835	0.7594
UserKNN	0.9382	0.7681	0.9296	0.7920
ItemKNN	0.9944	0.8242	0.9502	0.8160

Table 1: Rating prediction in original rating units; lower is better. RMSE and MAE are macro-averaged over test users. MatrixFormer-Impute uses global observed-training normalization; MatrixFormer retains each expert’s recipe as specified in the text. All six Cornac baselines are retained.

Dataset	Models	Tasks	Observed
MTEB	263	56	77.3%
MMLU	5,452	57	100.0%
Merged	118	114	31.1%
BenchPress	83	49	33.8%

Table 2: Released benchmark-score matrices used for evaluation. Observed entries are the available scores before holding out model rows.

Dataset	$k$	MatrixFormer	Gaussian
MTEB	5	0.7651	0.7498
	10	0.8261	0.8125
	15	0.8316	0.8259
MMLU	5	0.8992	0.9064
	10	0.8993	0.9118
	15	0.9155	0.9266

Table 3: Benchmark-score completion with greedy MI selection. Entries are mean

R^{2}

over ten folds at each task budget

k

, rounded to four decimals. All methods share context rows, selected tasks, target masks, and scoring rules. Numerical differences are descriptive.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Component	Setting
Value embedding $f_{\mathrm{val}}$	MLP $1\to 768\to 768$ with GELU
Missing entries	learned mask vector $m_{\mathrm{mask}}$ ; learned two-entry missingness embedding $e_{\mathrm{miss}}$ added to every cell
CLS border	$K_{c}=12$ row-CLS tokens $t^{\mathrm{row}}_{j}$ and $K_{r}=12$ column-CLS tokens $t^{\mathrm{col}}_{i}$ ; corner entry $(i,j)$ is $t^{\mathrm{col}}_{i}+t^{\mathrm{row}}_{j}$
Blocks	$L=8$ , width $d=768$ ; feature attention, datapoint attention, feed-forward
Attention	16 heads of dimension 48; rotary embeddings (base $10^{4}$ ) on 24 of the 48 dimensions; positions start after the border
Feed-forward	SwiGLU, $W_{3}\big(\mathrm{SiLU}(W_{1}h)\odot W_{2}h\big)$ , hidden width 1536

Appendix

Table 4: Shared architecture of MatrixFormer-Impute and MatrixFormer-RecSys.

Parameter	Role	Range
$s_{\max}$	Leading spectral scale	$[1,5]$
$\omega$	Initial head mixture weight	$[0,3]$
$h$	Stretched-exponential head scale	$[1,30]$
$p$	Head exponent	$[1,3]$
$H$	Head-to-tail transition scale	$[1,20]$
$\tau$	Power-law tail amplitude	$[0.2,1]$

Appendix

Table 5: Shared parameters of the three spectral prior families. All intervals specify uniform distributions in the raw generator’s units.

Explore similar work

Is One Layer Enough? Understanding Inference Dynamics in Tabular Foundation Models
May 7, 2026Amir Rezaei Balef, Mykhailo Koshil, Katharina EggenspergerTabular Foundation Models Transformer Architectures
TU Dortmund University, Dortmund, Germany · Lamarr Institute for Machine Learning and Artificial Intelligence, Dortmund, Germany · University of Tübingen, Tübingen, Germany
Data Language Models: A New Foundation Model Class for Tabular Data
May 7, 2026Eda Erol, Giuliano Pezzoli, Ozer Cem KelahmetTabular Learning Tabular Data
SchemaLabs
Auto-Fill: Learning to Predict Missing Values Accurately with Specialist Language Models
Jul 22, 2026Yurong Liu, Yeye He, Haoyu Dong +4Auto Research Table Reasoning Datasets
New York University · Microsoft Research

May 7, 2026, cs.LGJ/K move · Enter open · S save

Is One Layer Enough? Understanding Inference Dynamics in Tabular Foundation Models

Amir Rezaei Balef, Mykhailo Koshil, Katharina Eggensperger

TU Dortmund University, Dortmund, Germany · Lamarr Institute for Machine Learning and Artificial Intelligence, Dortmund, Germany · University of Tübingen, Tübingen, Germany

Tabular Foundation Models Transformer Architectures Layer-Wise Foundation Model Inference Dynamics Latent Dynamics

Open paper AskPDF arXiv

Transformer-based tabular foundation models (TFMs) dominate small to medium tabular predictive benchmark tasks, yet their inference mechanisms remain largely unexplored. We present the first large-scale mechanistic study of layerwise dynamics in 6 state-of-the-art tabular in-context learning models. We explore how predictions emerge across depth, identify distinct stages of inference and reveal latent-space dynamics that differ from those of language models. Our findings indicate substantial depthwise redundancy across multiple models, suggesting iterative refinement with overlapping computations during inference stages. Guided by these insights, we design a proof-of-concept, looped single-layer model that uses only 20% of the original model's parameters while achieving comparable performance. The code is available at https://github.com/amirbalef/is_one_layer_enough.

Abstract

Explore similar work

Is One Layer Enough? Understanding Inference Dynamics in Tabular Foundation Models

Data Language Models: A New Foundation Model Class for Tabular Data

Auto-Fill: Learning to Predict Missing Values Accurately with Specialist Language Models