README.md

---
license: apache-2.0
tags:
  - matrix-completion
  - missing-data-imputation
  - tabular
  - pytorch
---

# TabImpute and MatrixFormer

Pretrained transformers for completing numerical matrices with missing entries.
This repository contains the original TabImpute checkpoints and the two frozen
experts used by MatrixFormer+.

## Models

| Model | Checkpoint | Purpose |
|---|---|---|
| TabImpute | `tabimpute_001.ckpt` | Original entry-wise TabImpute model |
| TabImpute v2 | `tabimputev2.ckpt` | Earlier TabImpute v2 release |
| MatrixFormer-Impute | `matrixformer/impute/best_validation_checkpoint.pth` | General matrix completion; also called MatrixFormer in the comparison tables |
| Recommender specialist | `matrixformer/recsys/best_validation_checkpoint.pth` | Second frozen expert used by MatrixFormer+ |
| MatrixFormer+ | Both MatrixFormer checkpoints | Observed-error-fitted convex blend of the two experts |

Each MatrixFormer expert has 77,998,496 parameters and uses the Dwip v2-2
architecture (`new_stable`). Both were trained on synthetic matrices. The
checkpoints contain one network each, with no optimizer state. Their original
training configs are included alongside the weights.

MatrixFormer+ uses the observed entries of the supplied matrix to choose one
convex blending weight. The neural network weights stay frozen. This fitting
step does not use the true values of missing entries, and its observed-entry
reconstruction loss is an in-sample criterion rather than a held-out error
estimate. There is no separate MatrixFormer+ checkpoint.

## Installation

Use Python 3.10–3.12 and the published MatrixFormer source revision:

```bash
python -m pip install "tabimpute @ git+https://github.com/jacobf18/tabular.git@2cf54ac47dbcc1154b35e875b527ed0c42a9f416#subdirectory=mcpfn"
```

This installs the `tabimpute` package and Hugging Face download client. Use a
PyTorch build appropriate for your CUDA installation if running on a GPU.
The pinned source includes `StreamedImputePFN` and `TwoExpertRouter`; an older
PyPI release may lack these APIs. These are custom PyTorch checkpoints, rather
than Transformers `AutoModel` checkpoints.

## Run from Python

Supply a two-dimensional numeric NumPy array with `np.nan` for missing entries.
This example uses column normalization, predictive means, and the memory-bounded
streaming implementation. It preserves observed entries exactly.

```python
import json

import numpy as np
from huggingface_hub import hf_hub_download
from tabimpute.routing import TwoExpertRouter
from tabimpute.streamed_inference import StreamedImputePFN

repo_id = "Tabimpute/TabImpute"
revision = "main"
device = "cpu"


def load_expert(name):
    config_path = hf_hub_download(
        repo_id, f"matrixformer/{name}/config.json", revision=revision
    )
    checkpoint_path = hf_hub_download(
        repo_id,
        f"matrixformer/{name}/best_validation_checkpoint.pth",
        revision=revision,
    )
    with open(config_path, encoding="utf-8") as handle:
        model_config = json.load(handle)["model"]
    return StreamedImputePFN(
        checkpoint_path=checkpoint_path,
        json_config={"config": model_config},
        device=device,
        prediction_stat="mean",
        precision_mode="amp_bfloat16",
        max_num_rows=2000,
        max_num_cols=1000,
    )


matrix = np.array(
    [[1.0, 2.0, np.nan], [2.0, np.nan, 6.0], [3.0, 6.0, 9.0]],
    dtype=np.float64,
)
general = load_expert("impute")
completed = general.impute(matrix.copy())
print("MatrixFormer-Impute:", completed)

specialist = load_expert("recsys")
blend = TwoExpertRouter(
    general, specialist, policy="observed_blend", normalization="column"
)
completed_plus = blend.impute(matrix.copy())
print("MatrixFormer+:", completed_plus)
print("Blend decision:", blend.last_result.decisions["observed"])
```

Use `device = "cuda"` on a supported GPU. The published inference recipe uses
BF16 autocast (`amp_bfloat16`). `precision_mode="float32"` is an alternative
for hardware without BF16 support; results can differ with precision.
For reproducibility, replace `revision = "main"` with a Hub commit hash.

## Run from the command line

Download the example script, then run a small synthetic inference example:

```bash
hf download Tabimpute/TabImpute examples/run_matrixformer.py --local-dir tabimpute-models
python tabimpute-models/examples/run_matrixformer.py --model plus --device cpu
```

The script downloads only the selected experts and their metadata, verifies
SHA-256 hashes, and performs inference. It does not train a model. To run on your
own numerical matrix, save a two-dimensional `.npy` array containing `np.nan`
where values are missing:

```bash
python tabimpute-models/examples/run_matrixformer.py \
  --model impute --device cuda --input matrix.npy --output completed.npy

python tabimpute-models/examples/run_matrixformer.py \
  --model plus --device cuda --input matrix.npy --output completed_plus.npy
```

To download the complete MatrixFormer bundle once and work offline:

```bash
hf download Tabimpute/TabImpute --local-dir tabimpute-models \
  --include "matrixformer/*" "examples/run_matrixformer.py" "README.md"
python tabimpute-models/examples/run_matrixformer.py \
  --local-dir tabimpute-models --model plus --device cpu
```

The script defaults to `amp_bfloat16`, 2000-row by 1000-column context windows,
and CPU execution. Use `--precision float32` for full precision. Smaller
`--max-rows` and `--max-cols` values reduce context size; `0` uses the full axis.
Streaming reduces temporary allocations but the hidden matrix remains on the
chosen device. Large contexts still require substantial memory. CPU execution
is suitable for small examples; GPU execution is recommended for larger matrices.

## Preprocessing and evaluation scope

The examples above are a general-purpose column-normalized recipe. They are
not a reproduction of every table in the paper. Window size, orientation,
preprocessing, clipping, and precision can all change the predictions.

In particular, the reported recommender results use global normalization for
the general expert, two-way centering and user scaling for the specialist,
different context limits, rating-range clipping, and global normalization of
the blending loss. Loading the specialist with the generic example does not
reproduce that recipe. The frozen evaluation code and protocol descriptions are:

- [MatrixFormer+ protocols and comparisons](https://github.com/jacobf18/tabular/blob/2cf54ac47dbcc1154b35e875b527ed0c42a9f416/mcpfn/MATRIXFORMER_PLUS_RESULTS.md)
- [Observed-error blend and safeguards](https://github.com/jacobf18/tabular/blob/2cf54ac47dbcc1154b35e875b527ed0c42a9f416/mcpfn/OBSERVED_ROUTER.md)
- [Frozen evaluation entry point, including rating preprocessing](https://github.com/jacobf18/tabular/blob/2cf54ac47dbcc1154b35e875b527ed0c42a9f416/mcpfn/scripts/evaluate_two_expert_router.py)
- [Training history and checkpoint ancestry](https://github.com/jacobf18/tabular/blob/9c28633de0045bb51fc190a8941c12a8b7c3ac7d/AUTORESEARCH_TIMELINE.md)

Inputs must be numerical and finite apart from `np.nan`. Encode categorical
variables separately if needed. Entirely missing rows or columns have limited
observational support. Missingness patterns and matrix structures outside the
synthetic training distribution can reduce accuracy. The examples do not
provide calibrated uncertainty estimates or validate causal identification
assumptions for counterfactual applications.

## Provenance

`matrixformer/artifacts.json` records checkpoint and config hashes, original
run paths, and the compatible source revision. The general expert is the
selected large-matrix continuation from `large_matrices_98k_20260930`. The
specialist is `synthetic_user_scaled_parent_risk_400`. Config paths describe
the original research environment and are not required local directories for
the inference examples.

The original `tabimpute_001.ckpt` and `tabimputev2.ckpt` remain unchanged.
Calling `ImputePFN()` with its default arguments loads the original TabImpute
checkpoint, not MatrixFormer. Use the explicit loading recipe above for
MatrixFormer.

## Citation

For the original TabImpute method:

```bibtex
@article{feitelberg2025tabimpute,
  title={TabImpute: Accurate and Fast Zero-Shot Missing-Data Imputation with a Pre-Trained Transformer},
  author={Feitelberg, Jacob and Saha, Dwaipayan and Choi, Kyuseong and Ahmad, Zaid and Agarwal, Anish and Dwivedi, Raaz},
  journal={arXiv preprint arXiv:2510.02625},
  year={2025}
}
```