# N2: A Unified Python Package and Test Bench for Nearest Neighbor-Based Matrix Completion

arXiv: https://arxiv.org/abs/2506.04166 · HTML: https://mlsanity.com/paper/2506.04166/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 21 of its 46 references that are in the map.
- First posted on arXiv 4 Jun 2025.

## Builds on (most cited first, 21 of 21)

- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [Mistral 7B](https://mlsanity.com/paper/2310.06825/) (2023)
- [Gemma: Open Models Based on Gemini Research and Technology](https://mlsanity.com/paper/2403.08295/) (2024)
- [Kernel Mean Embedding of Distributions: A Review and Beyond](https://mlsanity.com/paper/1605.09522/) (2016)
- [Matrix Completion and Low-Rank SVD via Fast Alternating Least Squares](https://mlsanity.com/paper/1410.2596/) (2014)
- [Efficient multi-prompt evaluation of LLMs](https://mlsanity.com/paper/2405.17202/) (2024)
- [Personalized HeartSteps: A Reinforcement Learning Algorithm for Optimizing Physical Activity](https://mlsanity.com/paper/1909.03539/) (2019)
- [Causal Matrix Completion](https://mlsanity.com/paper/2109.15154/) (2021)
- [Missing Not at Random in Matrix Completion: The Effectiveness of Estimating Missingness Probabilities Under a Low Nuclear Norm Assumption](https://mlsanity.com/paper/1910.12774/) (2019)
- [Estimating Barycenters of Measures in High Dimensions](https://mlsanity.com/paper/2007.07105/) (2020)
- [Unitxt: Flexible, Shareable and Reusable Data Preparation and Evaluation for Generative AI](https://mlsanity.com/paper/2401.14019/) (2024)
- [Local nearest neighbour classification with applications to semi-supervised learning](https://mlsanity.com/paper/1704.00642/) (2017)
- [Design-based Analysis in Difference-In-Differences Settings with Staggered Adoption](https://mlsanity.com/paper/1808.05293/) (2018)
- [Matrix completion with data-dependent missingness probabilities](https://mlsanity.com/paper/2106.02290/) (2021)
- [Counterfactual inference in sequential experiments](https://mlsanity.com/paper/2202.06891/) (2022)
- [The Micro-Randomized Trial for Developing Digital Interventions: Experimental Design and Data Analysis Considerations](https://mlsanity.com/paper/2107.03544/) (2021)
- [Doubly robust nearest neighbors in factor models](https://mlsanity.com/paper/2211.14297/) (2022)
- [Benchmarking Generative AI for Scoring Medical Student Interviews in Objective Structured Clinical Examinations (OSCEs)](https://mlsanity.com/paper/2501.13957/) (2025)
- [Distributional Matrix Completion via Nearest Neighbors in the Wasserstein Space](https://mlsanity.com/paper/2410.13112/) (2024)
- [Learning Counterfactual Distributions via Kernel Nearest Neighbors](https://mlsanity.com/paper/2410.13381/) (2024)
- [Adaptively-weighted Nearest Neighbors for Matrix Completion](https://mlsanity.com/paper/2505.09612/) (2025)

## Abstract (arXiv)

Nearest neighbor (NN) methods have re-emerged as competitive tools for matrix completion, offering strong empirical performance and recent theoretical guarantees, including entry-wise error bounds, confidence intervals, and minimax optimality. Despite their simplicity, recent work has shown that NN approaches are robust to a range of missingness patterns and effective across diverse applications. This paper introduces N2, a unified Python package and testbed that consolidates a broad class of NN-based methods through a modular, extensible interface. Built for both researchers and practitioners, N2 supports rapid experimentation and benchmarking. Using this framework, we introduce a new NN variant that achieves state-of-the-art results in several settings. We also release a benchmark suite of real-world datasets, from healthcare and recommender systems to causal inference and LLM evaluation, designed to stress-test matrix completion methods beyond synthetic scenarios. Our experiments demonstrate that while classical methods excel on idealized data, NN-based techniques consistently outperform them in real-world settings.