Release frozen MatrixFormer experts and inference examples

7275749 verified 1 day ago

8.75 kB

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:

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.

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:

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:

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:

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:

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:

@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}
}