PaperlayerView exact revision on arXiv

Source-linked AI summary

MatrixFormer: A Foundation Model for Matrix Completion

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

arXiv:2610.06751v1cs.LGcs.AI

TL;DR

Matrix completion matters across settings where observed entries inform both rows and columns, but entrywise tabular approaches repeat context for each target. MatrixFormer uses a synthetic-pretrained matrix-native transformer, and the same model performs across imputation, recommendation, benchmark completion, and causal panel tasks.

  • Problem

    Matrix completion spans settings where observed entries inform both rows and columns, while entrywise tabular foundation models repeat observed context for every target.

    Source: §1 · p. 1

  • Method

    MatrixFormer represents cells as a two-dimensional token grid, alternates attention across matrix axes, and is pretrained on synthetic matrices to predict missing-entry distributions.

    Source: §1 · pp. 1–2, §2 · p. 2

  • Results

    With no fine-tuning, the same model has the lowest overall error on both imputation suites, lower rating error than collaborative-filtering baselines on two datasets, and the lowest normalized score in five of six causal-panel comparisons.

    Source: §4 · p. 11

  • Takeaways & Limitations

    The results support applying one fixed-weight completion model across tabular imputation, rating prediction, language-model benchmark completion, and causal panel data.

    Source: §1 · p. 2, §4 · p. 11

Abstract

from arXiv · show

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.

1 Introduction

MatrixFormer addresses matrix completion by replacing entrywise prediction, which repeats context, with a shared two-dimensional representation. It combines two fixed pretrained experts and achieves strong results across four completion domains.

1 Introduction

  • A shared two-dimensional representation avoids TabImpute’s expanded table of RC entry-level examples, which repeats observed values and raises memory and runtime costs for larger matrices.TabImpute’s entrywise approach turns every matrix entry into a supervised example containing its row and column context.Source: §1 · p. 1
  • MatrixFormer combines the general-purpose MatrixFormer-Impute with the recommender-specific MatrixFormer-RecSys, fitting one blend weight per input matrix while keeping both pretrained networks fixed.The weight minimizes normalized squared reconstruction error on observed entries, which are also part of the input; the blend is assessed on separately hidden entries and requires both experts’ forward passes.Source: §1 · pp. 1–2
  • Alternating attention across columns within rows and rows within columns lets each cell use same-unit and same-measurement information, while a shared decoder predicts all missing entries together.In the fixed-width runtime experiment on an NVIDIA H200, MatrixFormer completes a 1024 × 10 synthetic matrix in 0.0389 s versus 4.167 s for TabImpute, a 107× speedup.Source: §1 · p. 2
  • MatrixFormer-Impute is pretrained entirely on synthetic matrices spanning varied ranks, singular-value profiles, dimensions, and missingness patterns, including sparse rows and columns.At deployment, observed entries provide context while model weights remain fixed.Source: §1 · p. 2
  • MatrixFormer achieves state-of-the-art results across tabular imputation, rating prediction, language-model benchmark completion, and causal panel-data tasks with fixed model parameters.The evaluation includes dense and sparse observation patterns and matrices larger than those used in pretraining.Source: §1 · p. 2

2 Architecture and Training

MatrixFormer preserves a partially observed matrix as a two-dimensional grid of cell tokens, applying attention along its row and column axes before decoding missing entries. Its pretraining uses synthetic matrix families, varied masking and support patterns, teacher-distilled point predictions, and distributional losses.

2 Architecture and Training

  • MatrixFormer-Impute applies self-attention directly to a grid of cell tokens, preserving the matrix’s two axes rather than repeating row and column context in entry-level examples.The entrywise representation is used by TabImpute to make tabular prediction models applicable to completion.Source: §2 · p. 2
  • Observed cells receive a value embedding and an observation-status embedding, while missing cells receive a learned mask token and a distinct missing-status embedding.Zero is only a computational placeholder for missing values before the value embedding is replaced.Source: §2 · p. 2
  • Learned row- and column-CLS tokens border the grid, and each block alternates feature attention across columns within rows with datapoint attention across rows within columns.The two operations let a cell use other measurements for the same datapoint and the same variable across observations; RoPE is applied along the attended axis.Source: §2 · pp. 2–3, Figure 1 · p. 3
  • A shared decoder reads each contextualized cell independently to parameterize a predictive distribution, allowing all missing entries to be predicted in parallel in one forward pass.For point imputation, the predictive mean is used and observed entries are left unchanged.Source: §2 · pp. 3–4
  • Factoring attention across rows and columns reduces attention cost relative to full attention on the flattened grid while preserving row–column structure.With embedding width fixed, flattened dense self-attention requires O(R^2C^2) pairwise interactions per layer.Source: §2 · p. 4

2.1 Synthetic pre-training

  • Pretraining proceeds through base, main, and size-extension phases, with the final phase adapting the model to matrices as large as 2000 rows by 1000 columns.The main phase mixes low-rank, broader-rank spectral, and latent-factor families; its prior varies rank, spectral shape, singular-vector geometry, dimensions, and observed support.Source: §2.1 · p. 4
  • Episodes combine fixed-rate MCAR or mixed MCAR, MAR, and MNAR masking, and 75% additionally thin a random subset of rows or columns to model limited observation support.The mixed kernel includes structured patterns and nominal missing rates from 5% to 75%; scored entries are hidden entries, or hidden entries in thinned rows or columns.Source: §2.1 · p. 5
  • The training objective uses teacher predictions for the Huber point loss, while likelihood and CDF losses use the true hidden value; teachers are used only during training.The model is selected on fixed synthetic validation suites; downstream benchmarks were inspected during development, and the UCI context choice was an exploratory follow-up rather than blind validation.Source: §2.1 · pp. 5–6

3 Results

The results section evaluates MatrixFormer, a convex blend of MatrixFormer-Impute and MatrixFormer-RecSys, across causal inference, imputation, recommendation, and language-model evaluation. Reported results vary by task, including leading overall NRMSE on both imputation suites and higher benchmark-completion R2 on three of four score matrices at all three reported budgets.

3 Results

  • The evaluation covers causal inference, imputation, recommendation systems, and LLM evaluation, comparing MatrixFormer with standard benchmarks.MatrixFormer is introduced as a convex blend of MatrixFormer-Impute and MatrixFormer-RecSys.Source: §3 · p. 6

3.1 Inference runtime scaling

  • At 1024 rows and 10 columns, MatrixFormer's median runtime is 0.0389 s versus TabImpute's 4.167 s, a 107.0× speedup.The measured speedup increases from 2.1× at 16 rows to 107.0× at 1024 rows.Source: §3.1 · p. 6, Figure 2 · p. 6

3.2 Observed-error ensembling: MatrixFormer

  • MatrixFormer blends predictions from the general-purpose MatrixFormer-Impute and recommender-specific MatrixFormer-RecSys experts using one input-dependent weight per matrix.The weight is fitted on observed-cell prediction error, then used to combine missing-cell predictions; observed entries remain unchanged and both experts retain their pretrained parameters.Source: §3.2 · pp. 6–7

3.3 Causal panel data: counterfactual completion

  • MatrixFormer has the lowest normalized score in five of six causal-panel panels; HyperImpute is lowest for Kansas effect error, 0.119 against 0.125.Nonoverlapping 95% intervals mark clear treated-block RMSE leads on California smoking (0.159 against 0.212 for SDiD) and CPS (0.177 against 0.215); the other panels are within intervals.Source: §3.3 · p. 7, Figure 3 · p. 8

3.4 Tabular imputation: MissBench and UCI

  • 1.580 on MissBench narrowly beats TabImpute (legacy) at 1.585, and 0.998 on UCI beats ICE at 1.172 in overall NRMSE.MatrixFormer is also lowest on all three UCI mechanisms; on MissBench it is lowest for MNAR, while TabImpute leads MCAR and K-NN leads MAR.Source: §3.4 · p. 8, Figure 4 · p. 9, Figure 5 · p. 9

3.5 Recommender-system rating prediction

  • 0.8874/0.7279 RMSE/MAE on MovieLens 100K and 0.8617/0.7330 on Netflix small are lower than every displayed baseline.The Netflix RMSE margin over SVD (0.8621) is small and is not a significance claim.Source: §3.5 · p. 10, Table 1 · p. 10

3.6 Benchmark-score completion for language-model evaluation

  • At k = 15, MatrixFormer reaches R2 0.6464 against 0.3625 on Merged and 0.4687 against −0.0253 on BenchPress, but Gaussian leads MMLU at 0.9266 against 0.9155.MatrixFormer has higher R2 on MTEB, Merged, and BenchPress at all three reported budgets; across 180 settings it improves in 119 and is worse in 61, counts that summarize correlated settings rather than independent significance tests.Source: §3.6 · p. 11, Table 3 · p. 11

4 Conclusion

MatrixFormer preserves a partially observed matrix as a grid of cell tokens, alternates attention over rows and columns, and is pretrained only on synthetic matrices. Without fine-tuning, it achieves strong results across imputation, recommendation, benchmark-score completion, and causal panel data.

4 Conclusion

  • MatrixFormer has the lowest overall error on both imputation suites and the lowest normalized score in five of six causal-panel comparisons.It also has lower rating error than collaborative-filtering baselines on MovieLens 100K and Netflix small, and higher R2 than the Gaussian estimator on three of four benchmark-score matrices.Source: §4 · p. 11
  • A single forward pass completes a 1024 × 10 matrix 107× faster than TabImpute; MatrixFormer uses two such passes.Source: §4 · p. 11
  • The same model, without fine-tuning, was applied to tabular imputation, rating prediction, language-model benchmark-score completion, and causal panel data.Source: §4 · p. 11
  • Future work includes evaluating predictive distributions, extending context to larger matrices, and choosing ensemble weights on held-out entries.Source: §4 · p. 11

A.1 Architecture

MatrixFormer-Impute and MatrixFormer-RecSys share an architecture, while MatrixFormer combines their predictions. Table 4 lists the architecture settings.

A.1 Architecture

  • The two experts differ in their weights despite sharing the same architecture.Source: §A.1 · p. 14
  • MatrixFormer combines the predictions of the two experts.Source: §A.1 · p. 14
  • MatrixFormer-Impute and MatrixFormer-RecSys share the architecture listed in Table 4.Source: §A.1 · p. 14, Table 4 · p. 15

A.2 Bar distribution

The bar distribution assigns probability across central bins and two tails. Its distribution family has no trainable parameters; the model’s parameters enter through the logits.

A.2 Bar distribution

  • The density uses half-normal tails outside the outer borders and a uniform density within each interior bin.Source: §A.2 · p. 15
  • Each tail’s median is set to one bin width.Source: §A.2 · p. 15
  • The distribution family has no trainable parameters; model parameters enter only through the logits.Source: §A.2 · p. 15

A.3 Pre-training

Pre-training proceeds through a base phase, a diverse main phase, and a size-extension phase, with synthetic validation selecting checkpoints. The main phase uses teachers trained on synthetic tasks.

A.3 Pre-training

  • Two teachers are selected from earlier 64k-update runs, trained with true values in place of predicted means on different synthetic matrix distributions.Source: §A.3 · p. 15
  • The base phase trains on rank 1–11 latent-factor matrices with 10–50 rows, 5–50 columns, and 40% MCAR missingness for 17k updates.Source: §A.3 · p. 15
  • The main phase samples eight equally weighted components, combining matrix families with MCAR or mixed masking that includes MAR, MNAR, and structured patterns.It also includes spectral and latent-factor matrices, varied singular-vector structures, and sparse-support episodes; latent-factor matrices serve as rehearsal because spectral-only training degraded their performance.Source: §A.3 · p. 15
  • The size-extension phase runs 5k updates on progressively larger matrices and selects a checkpoint using synthetic panels, limiting error growth on main-phase-sized matrices to at most 3%.The selected checkpoint follows 4k updates; the main-phase checkpoint was selected after 98k updates because a 64k run was still improving at its end.Source: §A.3 · p. 16

A.4 Synthetic matrix generation

Synthetic matrices combine varied spectral and latent-factor priors with several singular-vector geometries. The constructions apply to both matrix axes without changing the model architecture.

A.4 Synthetic matrix generation

  • Shared spectral parameters are sampled independently and uniformly per matrix, and the training target retains all k signal components.The coefficient sequence need not be monotone; its decreasing rearrangement gives the singular spectrum because U and V have orthonormal columns.Source: §A.4 · p. 16, Table 5 · p. 16
  • Each axis independently selects Gaussian, coherent, smooth, or offset-bearing features with equal probability, then uses a thin QR factor for its singular vectors.The smooth construction adds a small Gaussian term to avoid rank-deficient feature matrices.Source: §A.4 · p. 16
  • The singular-vector feature construction applies to both rows and columns and does not alter the model architecture.The Gaussian construction yields a Haar-distributed orthonormal basis.Source: §A.4 · p. 16
  • Each latent-factor vector draws coordinates independently from a mixture of Gaussian, Laplace, Student-t, and other specified choices, while a fifth choice draws the vector jointly from a centered Dirichlet distribution.The Student-t degrees of freedom are uniform on {2, . . . , 10}.Source: §A.4 · p. 16
  • The latent dimension only upper-bounds realized rank, and broader-rank families may still draw small ranks because their names describe tail exponents.Rehearsal retains 10–50 rows and 5–50 columns during size extension.Source: §A.4 · pp. 16–17

B Causal-panel experiment details

The causal-panel benchmarks use placebo or semi-synthetic designs built from real panels, plus real California smoking outcomes with artificial treatment. They evaluate completion error and treatment-effect bias under zero true effects, using draw-normalized scores and draw-to-draw uncertainty intervals.

B Causal-panel experiment details

  • All three benchmarks have zero true treatment effects because hidden cells are untreated outcomes.CPS and Kansas simulate panels from fitted real-data models; California smoking uses real outcomes and artificial treatment.Source: B Causal-panel experiment details · p. 17
  • CPS placebo panels simulate state-year outcomes from rank-4 decompositions plus AR(2) residuals, then hide treated states in the final 10 years.The source panel has 50 states and 40 years, and each of nine outcome–assignment combinations has 200 draws.Source: B Causal-panel experiment details · p. 17
  • Kansas panels resample unit effects, factor loadings, and residuals around a fitted three-factor model, then hide the 16 post-treatment quarters for one sampled state.The source is quarterly log state GDP per capita for 50 states across 105 quarters.Source: B Causal-panel experiment details · p. 17
  • California smoking uses real cigarette-sales outcomes for 38 donor states and hides cells under simultaneous or staggered designs across five pre-treatment lengths.There are 100 draws per pre-treatment length.Source: B Causal-panel experiment details · p. 17
  • Within each draw, metrics are min-max normalized across seven main methods, then averaged over draws and equally weighted across designs.The true effect is zero, so absolute treatment-effect error equals absolute mean bias; 95% normal intervals use draw-to-draw variation, and shared source data may make them somewhat optimistic.Source: B Causal-panel experiment details · p. 17
Loading 2610.06751v1…