Can a matrix-native transformer complete missing entries in one forward pass?

One forward pass replaces entry-by-entry prediction

Earlier tabular foundation models treat matrix completion as entry-by-entry prediction: each cell becomes a supervised example carrying its row and column context, so an R×C matrix expands into a table of R×C examples whose features repeat the same observed values. That expansion raises memory and runtime cost as matrices grow. MatrixFormer instead represents the partially observed matrix as a grid of cell tokens with explicit missingness indicators, and attention alternates across the columns of each row and the rows of each column. A cell can therefore combine information about other measurements of the same unit with information about the same measurement across units.

The payoff is measured in the fixed-width runtime experiment on an NVIDIA H200: with width fixed at 10 columns and rows increasing geometrically from 16 to 1024, the speedup over the legacy entry-wise TabImpute grows from 2.1× at 16 rows to 107.0× at 1024 rows, where median runtime is 4.167 s for TabImpute and 0.0389 s for MatrixFormer. That comparison is a runtime measurement on one synthetic 1024×10 matrix under 40% MCAR, not an accuracy claim, and MatrixFormer itself uses two such passes because it ensembles two experts.

Synthetic pretraining keeps the weights fixed at deployment

MatrixFormer-Impute is trained entirely on synthetic matrices with varied ranks, singular-value profiles, dimensions, and missingness patterns, including sparse rows and columns that model limited observation support. At deployment, observed entries provide context and the model weights remain fixed. Training uses a fresh AdamW optimizer with weight decay 0.02 and gradient-norm clipping at 0.5, running 128k updates on batches of two matrices in FP32 parameters and BF16 computation, with a learning rate that warms up over 600 updates to 10⁻⁵ and then cosine-decays to 5% of that peak. A size-extension phase runs 5k updates at half that peak learning rate, enlarging matrices in stages.

Two details matter for reading the results. First, model selection uses synthetic data only, but the downstream benchmarks were inspected during development, and the UCI context choice is described by the authors as an exploratory follow-up rather than a blind validation. Second, the released model is an observed-error ensemble of two fixed pretrained experts: the general-purpose MatrixFormer-Impute and the recommender-specific MatrixFormer-RecSys. One mixture weight per input matrix minimizes the blend's normalized squared reconstruction error on observed entries, the two-model least-squares solution is analytic and clipped to [0,1], and the same weight combines all missing-cell predictions in that matrix while observations remain unchanged.

Causal panel data tests counterfactual completion

Estimating a treatment effect from panel data amounts to predicting the untreated outcomes of treated units in post-treatment periods from the remaining entries. The paper evaluates this on saved placebo benchmarks where the untreated values are known, so the true effect is zero: the CPS placebo study, the Kansas semi-synthetic factor design on one treated state, and the California-smoking donor pool with simultaneous or staggered pseudo-treatment. Two metrics are reported: block RMSE on the hidden treated block, and the absolute error of the implied average treatment effect on the treated. Because raw errors are not comparable across designs with different outcome scales, each metric is min-max normalized across the seven methods within each draw, averaged over draws, and designs are weighted equally.

MatrixFormer has the lowest normalized score in five of the six panels; HyperImpute is lowest for the Kansas effect error (0.119 against 0.125). The lead over the next-best method is clear, with non-overlapping 95% intervals, for treated-block RMSE on California smoking (0.159 against 0.212 for SDiD) and CPS (0.177 against 0.215); the other panels are within the intervals. Against HyperImpute alone, MatrixFormer is clearly ahead in three panels and tied in the other three. These are normalized scores over draws, so the ordering is the result, not the raw error magnitudes.

Imputation gains are narrow on MissBench and wider on UCI

Two imputation suites are used. MissBench is the OpenML-based suite used with TabImpute: 42 datasets under eleven missingness mechanisms at the nominal 40% setting, for 462 cases. The local UCI-style suite contains 12 datasets under MCAR, MAR and MNAR at 10%, 20%, 30%, 40% and 50% nominal missingness with ten repeats per setting, for 1,800 cases. Neither suite has a row train/test split: all rows and columns are retained, observed entries provide context, and predictions are scored only on masked cells. The metric is columnwise normalized RMSE, where each column's RMSE on its masked entries is divided by the standard deviation of its masked ground truth, averaged over valid columns of a case; within each missingness mechanism, repeats and rates are averaged within a dataset, datasets then receive equal weight, and the two suites are not pooled.

MatrixFormer has the lowest overall NRMSE on both suites: 1.580 on MissBench, narrowly ahead of legacy TabImpute at 1.585, and 0.998 on UCI, ahead of ICE at 1.172. On UCI it is also lowest for each of MCAR, MAR and MNAR. On MissBench it is lowest for MNAR, while legacy TabImpute is lowest for MCAR (0.819 against 0.866) and K-NN for MAR (1.075 against 1.086). The MissBench margin over TabImpute is therefore a near-tie in overall NRMSE, not a decisive gap. On UCI, MatrixFormer uses sequential windows of at most 2,000 rows and 1,000 columns, the largest matrices seen in pre-training, with observed-column normalization computed once from the entire incomplete matrix and overlapping tail-window predictions averaged; only rows are split, since these datasets have at most 90 columns. The authors state this single training-derived window limit was adopted after inspecting whole-matrix results, so it is an exploratory protocol choice.

Benchmark-score completion beats a Gaussian estimator on sparse matrices

Evaluating a language model on every benchmark is expensive, so the paper asks whether MatrixFormer can reconstruct a model's unmeasured benchmark scores after observing a small subset of its evaluations. Rows are models and columns are benchmark tasks; this is score-matrix completion, not a new evaluation of generated answers. Four fixed snapshots are used: MTEB (263 models, 56 tasks, 77.3% observed), MMLU (5,452 models, 57 tasks, 100.0% observed), Merged (118 models, 114 tasks, 31.1% observed), and BenchPress (83 models, 49 tasks, 33.8% observed). Each matrix is split into ten model-row folds with seed 42, holding out about 10% of models per fold; for each held-out model its available scores among k selected tasks are exposed and its originally observed, unselected scores are predicted. Naturally missing entries are not ground truth, so on sparse matrices k selected tasks may expose fewer than k scores.

The metric is R² on held-out targets, averaged over the ten folds. Under greedy mutual-information selection at k ∈ {5, 10, 15}, MatrixFormer has the higher R² on MTEB, Merged and BenchPress at all three budgets, most clearly on the two sparsest matrices: 0.6464 against 0.3625 on Merged and 0.4687 against −0.0253 on BenchPress at k = 15. The Gaussian conditional-mean estimator is higher on MMLU (0.9266 against 0.9155 at k = 15). Across all 180 dataset/selection/budget settings, MatrixFormer improves on the Gaussian baseline in 119 and is worse in 61; the paper notes these counts summarize correlated settings rather than independent significance tests.

Rating prediction is macro-averaged over users

Rating completion is evaluated on MovieLens 100K and the Netflix small subset distributed by Cornac, sparse user–item matrices with ratings on a 1–5 scale. The benchmark's frozen 80/20 rating-entry holdout is reused, scoring 19,971 ratings on MovieLens 100K and 121,470 on Netflix small; this is an entry holdout, not a held-out-user split. Following the Cornac harness, RMSE and MAE are computed within each test user and then averaged equally over users, so the reported RMSE is the mean of per-user RMSEs rather than a pooled rating-level RMSE.

MatrixFormer obtains RMSE/MAE of 0.8874/0.7279 on MovieLens 100K and 0.8617/0.7330 on Netflix small, lower than every displayed baseline. The Netflix RMSE margin over SVD (0.8621) is small and the paper explicitly does not make a significance claim for it. One baseline caveat is stated directly: BPR is trained with a pairwise ranking objective rather than for rating calibration, so its rating errors do not measure its ranking quality. MatrixFormer's predictions are clipped to the observed training rating range.

The zero-shot claim holds unevenly across tasks

Read together, the results do not support a uniform "best everywhere" reading. The clearest advantages are the runtime scaling on the 1024×10 synthetic matrix, the treated-block RMSE on California smoking and CPS where 95% intervals do not overlap the next-best method, and the two sparsest benchmark-score matrices at k = 15. The narrowest margins are the MissBench overall NRMSE (1.580 against 1.585 for legacy TabImpute) and the Netflix RMSE over SVD (0.8617 against 0.8621), neither of which the paper presents as a significance claim. MMLU is the one benchmark-score matrix where the Gaussian estimator is higher at every reported budget.

The most useful distinction for a reader is between the tasks where MatrixFormer's advantage is large and interval-separated and the tasks where it is a near-tie or a loss. The paper's own framing is that these results position MatrixFormer as a general-purpose foundation model for matrix completion, with a single forward pass completing a 1024×10 matrix 107× faster than TabImpute while MatrixFormer itself uses two such passes. The authors list evaluating the predictive distributions, for example as intervals for treatment effects, extending the context to larger matrices, and choosing the ensemble weight on held-out entries as directions for future work, which marks the boundary of what the current zero-shot evidence covers.

Sources used in this answer