A pretrained transformer that fills in the missing entries of any numerical matrix: ratings, panels, sensor grids, tables. No training, no fine-tuning, no hyperparameters to tune.
Pretrained once on synthetic matrices. At inference, one forward pass completes your matrix. No weights are updated.
Attention along rows and columns works for any number of rows and columns. Large matrices are processed in overlapping windows.
Only missing cells are predicted. Every observed value is returned exactly, in its original units.
One set of frozen weights, with no fine-tuning, evaluated across four domains: tabular imputation, causal inference on panel data, recommender systems and language-model evaluation. All numbers are from the paper. Hover any chart for exact values.
Two suites, scored on masked cells with column-normalized RMSE (NRMSE). MissBench: 42 OpenML datasets under 11 missingness mechanisms at 40% missingness. UCI: 12 datasets under MCAR, MAR and MNAR at 10–50% missingness, with 10 repeats each. MatrixFormer has the lowest overall error on both.
Lower is better. Whiskers are 95% intervals from 10,000 dataset-bootstrap resamples. Rows are sorted by overall NRMSE, and that order is kept across patterns. † marks incomplete coverage.
Estimating a treatment effect means predicting what treated units would have done untreated. On placebo panels where the true effect is zero, we hide the treated block and score its completion. Benchmarks: CPS state wage panels, the Kansas semi-synthetic design, and the California smoking donor pool. Scores are min–max normalized within each draw across the seven methods: 0 is best and 1 is worst.
Lower is better. Means over draws, with the designs within a benchmark weighted equally. MatrixFormer is lowest in five of six panels; HyperImpute is lowest for the Kansas effect error (0.119 vs 0.125).
Rating prediction on sparse user–item matrices, with a frozen 80/20 rating holdout. Baselines are Cornac's collaborative-filtering models. RMSE and MAE are in rating units, averaged over test users. MatrixFormer is lowest on both datasets and both metrics; on Netflix, its RMSE margin over SVD is small.
Lower is better. The axis is zoomed to separate the top methods. BPR is trained for ranking, not rating calibration, so it falls off the scale; its value is shown at the edge.
Rows are language models and columns are benchmark tasks. For a held-out model we reveal its scores on k tasks chosen by greedy mutual information, then predict the rest. The baseline is the Gaussian conditional-mean estimator from the task-selection release. MatrixFormer is higher on MTEB, Merged and BenchPress, with the largest gains on the two sparsest matrices. The Gaussian estimator is slightly higher on dense MMLU.
R² on held-out scores, higher is better, averaged over ten model folds. Each panel has its own y-axis. Observed fractions: MTEB 77%, MMLU 100%, Merged 31%, BenchPress 34%.
TabImpute builds a separate prediction problem for every missing entry. MatrixFormer completes the whole matrix in one forward pass. With 10 columns, the speedup grows from 2.1× at 16 rows to 107× at 1,024 rows: 4.17 s for TabImpute vs 0.039 s for MatrixFormer-Impute.
Median runtime on one NVIDIA H200, 10 columns, both axes logarithmic.
The paper's benchmark runs use task-specific preprocessing that this package does not include (for example, bias centering for ratings), so running the package's generic recipe will not reproduce every number exactly. See Good to know.
Each matrix cell becomes a token. Eight transformer blocks alternate attention along rows and along columns, so every cell can borrow information from its row and its column. A decoder then predicts a distribution over the missing value.
One expert network, about 78M parameters. MatrixFormer+ runs two of these experts and blends them (below).
Observed values are standardized per column, so columns in very different units become comparable. Predictions are mapped back afterwards.
Observed values pass through a small MLP. Missing cells get a learned mask token. A missingness embedding tells the model which is which, and learned CLS tokens border each row and column.
Each block first lets cells in the same row attend to each other, then cells in the same column, then applies a per-cell feed-forward layer. Stacking eight blocks spreads information across the whole matrix.
For each cell the decoder outputs probabilities over 5000 value bins. The prediction is the mean of that distribution, which is written into the missing cells only.
Everything below is read directly from the inference code in src/matrixformer and the released expert configs.
Each expert has exactly 77,998,496 parameters. The eight transformer blocks hold 85% of them. The decoder is large too, because its last layer maps 1536 features to 5000 bin logits. The cell encoder is tiny: it only has to lift one scalar into a 768-dimensional token.
| Component | Parameters | Share | What it contains |
|---|---|---|---|
| Cell encoder | 612,864 | 0.8% | Value MLP (1 → 768 → 768), mask token, 2 missingness embeddings, 12 + 12 CLS embeddings |
| Each block | 8,269,827 | 10.6% | Two attention layers (2.36M each), SwiGLU feed-forward (3.54M), 3 RMSNorms, 3 residual scales |
| All 8 blocks | 66,158,616 | 84.8% | |
| Decoder | 11,227,016 | 14.4% | MLP 768 → 1536 → 1536 → 5000 with GELU |
An R × C input becomes an (R + 12) × (C + 12) grid of 768-dimensional tokens. There are no position embeddings at this stage; position enters later through RoPE inside attention.
Linear(1, 768) → GELU → Linear(768, 768).CLS tokens act as learned scratch registers. The decoder never reads them: it decodes only the R × C data cells.
Full attention over every cell would compare all RC tokens with each other. MatrixFormer instead factorizes attention along the two axes. Each block runs two separate attention layers with their own weights: first every row attends across its columns, then every column attends across its rows. After those two steps, a cell has seen its row, and through each cell in its column, every other row too. A single block already has a receptive field covering the whole matrix. Eight blocks repeat this mixing.
| Cost per block, 2000 × 1000 window | |
|---|---|
| Full attention over all tokens | 4.1 × 10¹² pairs |
| Axial (rows, then columns) | 6.2 × 10⁹ pairs |
| Reduction | ≈ 670× |
Token pairs scored by attention, including the 12 CLS rows and columns. Axial cost grows as RC(R + C) instead of (RC)².
scaled_dot_product_attention).For each data cell, a three-layer GELU MLP outputs 5000 logits. A softmax turns them into probabilities over 5000 value bins defined by fixed borders (assets/borders.pt), in units of column standard deviations. The imputed value is the mean of that distribution:
Bin borders per one-standard-deviation interval, from −16 to +16. 37 borders lie outside this range.
Means and population standard deviations come from the observed entries of the whole matrix before any windowing. An entirely missing column gets mean 0 and scale 1.
Each axis is cut into windows of at most 2000 rows and 1000 columns. The last window is shifted back to end at the edge, so it overlaps its neighbor. Cells covered twice get the average of both predictions. For example, 4500 rows split into [0, 2000), [2000, 4000) and [2500, 4500).
Within a window, attention runs in batches of 32 rows or columns, the feed-forward in chunks of 8192 tokens, with in-place residual updates. Every sequence is still attended in full, so streaming bounds memory without changing the math.
The hidden grid stays in float32 even under BF16 autocast: about 6.3 GB for a full 2000 × 1000 window, on top of the weights. Use smaller max_rows and max_cols to reduce it, at the cost of less context per window.
Like prior-data fitted networks (TabPFN) and TabImpute, MatrixFormer never sees a real dataset during training. It learns from an endless stream of synthetic complete matrices, each masked into a reconstruction task. Details below follow Section 2 and Appendix A of the paper.
Each phase continues from the best checkpoint of the previous one. Two frozen teacher models, trained beforehand from the base checkpoint, guide the main and size-extension phases.
Size extension plays the role of context-length extension in language models. Attention cost grows with matrix size, so most training uses moderate sizes. A short final phase then adapts the model, including the range of its rotary positions, to large matrices.
| Size-extension stage | Updates | Rows | Columns |
|---|---|---|---|
| Bridge | 1–500 | 256–512 | 128–256 |
| Large | 501–1500 | 512–1024 | 256–512 |
| Larger | 1501–3000 | 768–1536 | 384–768 |
| Full | 3001–5000 | 1024–2000 | 512–1000 |
The main phase draws from four equally likely matrix families. Together they vary rank, the shape of the singular-value spectrum, the geometry of the singular vectors, matrix size, and which entries are observed.
| Family | Rank k | Tail exponent α |
|---|---|---|
| Low-rank | 2–16 | 0.2–2 |
| Slow tail | 2–160 | 0–0.5 |
| Steep tail | 2–160 | 1–2.5 |
| Latent factor | 1–11 | n/a |
The three spectral families build T = U diag(s) Vᵀ with orthonormal U, V. Each singular value blends a fast-decaying head with a power-law tail and a terminal taper:
The input adds tiny Gaussian noise (σ ≤ 10⁻³), but the target keeps the full clean signal. The latent-factor family sets T = PQᵀ, with each factor vector drawn from a Gaussian, Laplace, Student-t, spike-and-slab or centered Dirichlet distribution. It is kept at base-phase sizes as rehearsal, because training on spectral matrices alone degraded performance on it.
Drawn live in your browser from the formula above with the paper's parameter ranges (Table 7), two draws per family. Rank is capped at 100 for display; each curve drops to zero after its rank k.
Each axis independently picks one of four constructions before QR orthonormalization: Gaussian (a uniformly random basis), coherent (variation concentrated in a few rows), smooth (random cosine features) or offset-bearing (a near-constant first direction).
Each family is paired with one of two masking schemes, giving eight equally weighted components: fixed 40% MCAR, or a mixed kernel that samples a missingness mechanism per matrix. Missingness can therefore depend on the values themselves.
| Mixed mechanism | Probability |
|---|---|
| MCAR: missing completely at random | 0.30 |
| MAR: depends on observed values | 0.20 |
| MNAR: depends on the hidden values | 0.20 |
| Censoring below or above a quantile | 0.10 |
| Two-phase subsampling | 0.10 |
| Skip logic | 0.05 |
| Cluster-level missingness | 0.05 |
The model is trained to predict a full distribution for every hidden entry. Its probabilistic terms are supervised by the true value. Its point prediction is instead distilled toward two frozen teachers: models with the same architecture that receive the same masked input. Their averaged predictive mean, with gradients stopped, becomes the point target.
| Setting | Main phase | Size extension |
|---|---|---|
| Optimizer | AdamW, β = (0.9, 0.999), weight decay 0.02 | same |
| Gradient clipping | 0.5 | 0.5 |
| Peak learning rate | 1 × 10⁻⁵ | 5 × 10⁻⁶ |
| Schedule | 600 warmup updates, cosine to 5% | 200 warmup updates, cosine to 5% |
| Batch | 2 matrices | 2 matrices, accumulated serially |
| Precision | FP32 parameters and Adam moments, BF16 compute | |
| Model selection | lowest macro RMSE on a fixed synthetic validation suite | best large-matrix error, with ≤ 3% loss on main-size matrices |
MatrixFormer-RecSys shares the architecture and differs only in its weights. It is fine-tuned on synthetic rating matrices that look like recommender data: very sparse, with heterogeneous users and popular items. The settings below come from the released checkpoint's config; the paper describes this expert only briefly.
MatrixFormer+ runs two experts with the same architecture: a general imputation model and a recommender-system specialist trained on sparse, rating-like matrices. It then picks one mixing weight for your matrix from how well each expert reconstructs the entries you did observe.
This is a closed-form fit of one number per matrix. No neural network weights change, and no held-out values are used.
Needs Python 3.10+, PyTorch 2.4+, NumPy and huggingface-hub. Weights download automatically from Hugging Face on first use. A GPU is optional.
git clone https://github.com/jacobf18/MatrixFormer.git
cd MatrixFormer
python3 -m venv .venv && source .venv/bin/activate
python -m pip install "torch>=2.4,<3" --index-url https://download.pytorch.org/whl/cpu
python -m pip install .
python -m matrixformer --device cpu --precision float32 # runs a small demo
For an NVIDIA GPU, install PyTorch from the PyTorch selector instead. See the README for Windows and troubleshooting.
import numpy as np from matrixformer import MatrixFormerPlus matrix = np.array([ [1.0, 2.0, np.nan], [2.0, np.nan, 6.0], [3.0, 6.0, 9.0], ]) model = MatrixFormerPlus.from_pretrained(device="cpu", precision="float32") completed = model.impute(matrix) print(completed) print(model.last_decision) # blend weight and reconstruction errors
Use device="cuda" on a GPU, where the default amp_bfloat16 precision is used. Reuse the model object across matrices.
python -m matrixformer --input matrix.npy --output completed.npy \
--device cpu --precision float32
Input is a .npy array with NaN for missing entries. Add --model impute or --model recsys to run a single expert, or --download-only --local-dir ./checkpoints to prepare an offline bundle.
| Model | Experts | Weights |
|---|---|---|
MatrixFormerPlusdefault · --model plus | General + specialist, blended | ~624 MB |
MatrixFormerImpute--model impute | General only | ~312 MB |
MatrixFormerRecSys--model recsys | Recommender specialist only (how it was trained) | ~312 MB |
| Parameters per expert | 77,998,496 |
| Blocks · width · heads | 8 · 768 · 16 |
| Feed-forward width | 1536 (SwiGLU) |
| Output bins | 5000 |
| Default context window | 2000 × 1000 |
np.nan. Zeros and sentinels like -999 are treated as observed.If you use MatrixFormer, please cite the paper.
@article{saha2026matrixformer,
title={MatrixFormer: A Foundation Model for Matrix Completion},
author={Saha, Dwaipayan and Feitelberg, Jacob and Choi, Kyuseong and Dwivedi, Raaz and Agarwal, Anish},
journal={arXiv preprint arXiv:2610.06751},
year={2026}
}