N squared: A Unified Python Package and Test Bench for Nearest Neighbor-Based Matrix Completion
cs.LG, stat.CO, stat.ML
Submitted: 2025-06-04
Updated: 2026-09-13
Comments: 22 pages, 6 figures
Code: https://github.com/aashish-khub/NearestNeighbors
License: http://creativecommons.org/licenses/by/4.0/
The gist: 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,
Terminology
Abstract
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 N squared, 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, N squared 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.
Sources
- Doubly Robust Inference in Causal Latent Factor Models
- Unitxt: Flexible, Shareable and Reusable Data Preparation and Evaluation for Generative AI
- Estimating Barycenters of Measures in High Dimensions
- Learning Counterfactual Distributions via Kernel Nearest Neighbors
- Counterfactual inference in sequential experiments
- Doubly robust nearest neighbors in factor models
- Distributional Matrix Completion via Nearest Neighbors in the Wasserstein Space
- Benchmarking Generative AI for Scoring Medical Student Interviews in Objective Structured Clinical Examinations (OSCEs)
- Measuring Massive Multitask Language Understanding
- Mistral 7B
- Two-Sided Nearest Neighbors: An adaptive and minimax optimal procedure for matrix completion
- Adaptively-weighted Nearest Neighbors for Matrix Completion
- Gemma: Open Models Based on Gemini Research and Technology
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: Local-Consistency Optimisation for Bayesian Neural Networks
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks