Doubly robust nearest neighbors in factor models

arXiv:2211.14297 · stat.ML, cs.LG · Submitted 2026-08-18 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Next we'll be talking about the paper "Doubly robust nearest neighbors in factor models".

Jane: The paper was written by Raaz Dwivedi, Caleb Chin, Sabina Tomkins, Predrag Klasnja, Susan Murphy et al. from Harvard University and Massachusetts Institute of Technology and University of Michigan.

Tom: Stay tuned as we take you through the paper and discuss its implications.

Jane: We also have Lu with us today — senior AI researcher at Tsinghua.

Tom: We also have Meng with us today — lead engineer at a mysterious AI startup.

Jane: We also have Lalam with us today — the in-house Large Language Model.

Tom: Alright, let's get started.

Title: Tom: Alright, welcome back to the show, everyone! Today we're digging into a fresh arXiv paper that's got a mouthful of a title: "Doubly Robust Nearest Neighbors in Factor Models." Jane, I gotta say, just reading that title made me think of two things — old-school recommendation algorithms and a statistical superpower.

Jane: Tom, you're not wrong on either count. Let's unpack it. The paper comes from a big team — folks at Harvard, MIT, and the University of Michigan. And the core problem is something we all bump into daily: missing data. Think of a streaming service where you know what some users watched, but not all of them, and you want to guess what a specific person would rate a specific movie.

Tom: Right, and that's the "factor model" part. The idea is that there's some hidden reason — like a user's taste and a movie's style — that determines the rating. But here's the catch: the authors are dealing with a non-linear relationship between those hidden factors, which is way more realistic than the simple linear stuff.

Jane: Exactly. And "nearest neighbors" is the classic fix. You find similar users or similar movies and average their ratings. But the paper's title says "doubly robust," which is the clever twist. It means the method works even if you only have good neighbors on one side — say, similar users but no similar movies.

Tom: So it's like having two flashlights in a dark room. If one dies, you still have light from the other. But if both work, you get a much clearer picture. That's the gist of the "doubly robust" part — it's not just a backup; it actually combines both sources to do better than either alone.

Jane: And that's a big deal for real-world applications, not just movies. The paper mentions digital health, like deciding when to send a reminder to a patient. You might have lots of similar patients but few similar time points, or vice versa. This method handles both gracefully.

Tom: So, the title is a promise: robust to missing neighbors, and robust in a way that improves accuracy. I'm already excited to see how they pull it off mathematically. What do you think the actual algorithm looks like?

Jane: Well, Tom, that's exactly what we're going to dig into next. The paper doesn't just wave its hands; it gives a concrete recipe. And trust me, the recipe is surprisingly intuitive once you see it.

Paper discussion segment 2: Tom: So, Jane, we've got the title and the big idea. Now let's talk about the actual method. The paper introduces something called the DR-NN estimate, and I have to say, the construction is almost elegant in its simplicity.

Jane: It really is. So, imagine you want to estimate the outcome for user *i* at time *t*. The unit-NN approach finds similar users and averages their outcomes at time *t*. The time-NN approach finds similar times and averages user *i*'s outcomes at those times. Each of these has a bias — they're approximations.

Tom: And the DR-NN estimate is like a clever correction. It takes the time-NN estimate, but then subtracts a bias term that's estimated using the unit neighbors. And symmetrically, it can be seen as correcting the unit-NN estimate using the time neighbors. The paper shows that this correction makes the bias much smaller.

Jane: Right. In the paper, they show that for a simple bilinear model, the error of the DR-NN estimate is roughly the *product* of the errors of the two vanilla approaches, not the sum. That's a huge difference. If each vanilla method has an error of, say, zero point one, the product is zero point zero one, which is a hundredfold improvement.

Tom: And that's not just a theoretical trick. The paper proves it with non-asymptotic bounds — that's math-speak for "guarantees that hold with high probability for finite data, not just when you have infinite data." They also show that if one of the vanilla methods is useless, the DR-NN still performs as well as the other one. That's the "doubly robust" property in action.

Jane: Exactly. And they don't stop there. They also provide a confidence interval for the estimate, which is crucial for decision-making. If you're a doctor deciding on a treatment, you don't just want a point estimate; you want to know how uncertain it is.

Tom: And the confidence interval is also narrower than the vanilla ones, again by roughly a quadratic factor. So you get better estimates and tighter uncertainty. That's a win-win.

Jane: But, Tom, I should mention that the paper does make some assumptions. The noise has to be bounded, and the missingness is completely at random. Those are standard in the field, but they're worth keeping in mind.

Tom: Sure, but the authors are pretty upfront about that. And they show that the method is still a big step forward. Now, I'm curious about the practical side. How does this actually perform in the examples they give?

Jane: That's the perfect segue to our next segment. They run through a couple of concrete examples — one with discrete factors and one with continuous factors — and the improvements are pretty striking.

Paper discussion segment 3: Tom: So, Jane, let's get into the concrete numbers. The paper gives two main examples. First, the discrete case, where users and times come from a finite set of types. Second, the continuous case, where the factors are drawn from a uniform distribution.

Jane: Right. In the discrete case, with a constant observation probability, the DR-NN error scales like M over N, where M is the number of types and N is the number of users. That's a parametric rate — it decays as fast as possible. The vanilla unit-NN, on the other hand, scales like one over the square root of N plus M over N. So the DR-NN is a quadratic improvement in the leading term.

Tom: And in the continuous case, it's even more interesting. The vanilla NN error scales like N to the power of negative two over (d+two), where d is the dimension of the latent factors. The DR-NN error scales like N to the power of negative four over (d+four). For a fixed dimension, that's a much faster decay.

Jane: Let's put that in plain terms. If d is, say, four the vanilla error decays like N to the minus one-third, while the DR-NN error decays like N to the minus one-half. That's the difference between needing a thousand users versus a hundred users to get the same accuracy.

Tom: And that's not just a theoretical curiosity. The paper also shows that the confidence intervals shrink accordingly. So you get better point estimates and tighter intervals, which is exactly what you want in applications like personalized medicine or recommender systems.

Jane: But, Tom, there's a subtlety. The paper shows that the DR-NN is never worse than the better of the two vanilla methods. But to get the full quadratic improvement, you need both unit and time neighbors to be reasonably good. If one is terrible, you still get the performance of the other.

Tom: That's the "doubly robust" property we talked about. It's like having a car with two engines. If one fails, you can still drive on the other. But if both work, you get twice the speed. And the paper proves this rigorously, not just with simulations.

Jane: And they also extend the analysis to non-linear factor models, which is a big deal. The bilinear case is nice, but real-world data is rarely that simple. The paper shows that the DR-NN still provides a significant improvement, though the gains are a bit more modest.

Tom: So, the takeaway is that this isn't just a tweak. It's a fundamentally better way to combine information from two axes of a matrix. And the math is solid. I'm really impressed by the depth of the analysis.

Jane: Me too. And it makes me wonder about the broader impact. We've got a method that's simple to implement, has strong guarantees, and works in realistic settings. That's a recipe for adoption.

Conclusion: Tom: Alright, Jane, let's wrap this up. We've been talking about "Doubly Robust Nearest Neighbors in Factor Models" from Harvard, MIT, and Michigan. And I think the big takeaway is that this isn't just a new estimator; it's a new way of thinking about robustness in matrix completion.

Jane: Absolutely. The paper shows that by combining unit and time neighbors in a clever way, you get an estimate that's never worse than the best of the two vanilla approaches, and often much better. The math is rigorous, the examples are concrete, and the potential applications are everywhere — from recommender systems to digital health.

Tom: And let's not forget the confidence intervals. They're tighter, which means better decision-making under uncertainty. That's the kind of thing that can actually change practice.

Jane: Right. And while the paper makes some standard assumptions, like missingness at random, the core idea is robust enough that it's likely to inspire follow-up work in more complex settings, like adaptive experiments or causal panel data.

Tom: So, we're saying goodbye to this paper, but not to the ideas. I'm already thinking about how this could be extended to tensors or to settings with side information.

Jane: Definitely. But for now, we've got a solid understanding of a really clever piece of work. Thanks for joining us, everyone. We'll be back soon with the next paper.

Tom: Until then, keep your neighbors close, and your doubly robust ones closer. See you next time!

Raaz Dwivedi, Caleb Chin, Sabina Tomkins, Predrag Klasnja, Susan Murphy, Devavrat Shah

Harvard University · Massachusetts Institute of Technology · University of Michigan

stat.ML, cs.LG

Submitted: 2026-08-18

Updated: 2026-08-19

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 63/100

The gist: The paper addresses matrix completion with missing data under a latent factor model.

Key concepts

Factor Model: A model that uses hidden factors, such as a user's taste or a movie's style, to determine outcomes and predict missing data. The paper explores these models with non-linear relationships between factors, which is more realistic for real-world applications like digital health or recommendation systems.
Doubly Robust: A property where a method works effectively even if only one of two different approaches is accurate. The DR-NN method is robust to missing neighbors; if one source of information is unreliable, the method still performs as well as the other, but improves significantly if both are good.
DR-NN Estimate: A method that corrects bias by combining the unit-NN approach, which averages outcomes from similar users, with the time-NN approach, which averages outcomes from similar time points. This combination results in an error rate that is much smaller than using either approach individually.

Terminology

Matrix Completion — 7× in this paper · explained in 3 episodes
Loss Function — 5× in this paper · explained in 2 episodes
Robustness — 4× in this paper · explained in 5 episodes
Consistency — 3× in this paper · explained in 3 episodes
Accuracy — 2× in this paper · explained in 1 episode
Core — 2× in this paper · explained in 2 episodes
Curse of Dimensionality — 2× in this paper · explained in 6 episodes
Generalization — 2× in this paper · explained in 13 episodes
Interpretability — 2× in this paper · explained in 15 episodes
Adaptation — 1× in this paper · explained in 1 episode
Adaptive Algorithm — 1× in this paper · explained in 1 episode
Distillation — 1× in this paper · explained in 2 episodes
Knowledge Distillation — 1× in this paper · explained in 11 episodes
Minimax optimality — 1× in this paper · explained in 2 episodes
Non-asymptotic Analysis — 1× in this paper · explained in 2 episodes
alpha — 1× in this paper · explained in 1 episode

Summary

The paper addresses matrix completion with missing data under a latent factor model. The observed data follows:


Yi,t = θi,t + εi,t if Ai,t = 1, and Yi,t = ⋆ if Ai,t = 0

where θi,t is the mean outcome for unit i at time t, εi,t is mean-zero noise with variance σ2, and Ai,t is a binary indicator of observation. The mean parameters admit a (possibly non-linear) factorization:


θi,t = f(ui, vt)

for latent factors ui (unit factors) and vt (time factors), where f is an unknown function. The goal is entry-wise estimation of θi,t.

The paper introduces a doubly robust nearest neighbor (DR-NN) estimator that combines unit-NN and time-NN approaches. The authors state: We present a new NN estimate that brings together the unit and time-NN estimates to provide a significant improvement for providing entry-wise guarantee for estimating Θ. The estimator is doubly robust in two ways: (1) it provides consistent estimates as long as either good row or good column neighbors exist, and (2) when both good row and column neighbors exist, it provides a near-quadratic improvement in the non-asymptotic error and a significantly narrower asymptotic confidence interval compared to both vanilla NN approaches.

Unit-NN: Given threshold η1, define unit neighbors:


S unit i,t,η1 = j ≠ i: ρ̂ unit i,t (j) ≤ η1

where ρ̂ unit i,t (j) is the mean squared distance of observed outcomes for units i and j at time points except t. The unit-NN estimate is the average of observed outcomes of these neighbors at time t.

Time-NN: Analogously, define time neighbors S time i,t,η2 by comparing outcomes across units except i, and the estimate averages observed outcomes of unit i across these time neighbors.

DR-NN: The doubly robust estimate is defined as:


θ̂ DR i,t,η = [Σ j∈S unit, t'∈S time Ai,t' Aj,t Aj,t' (Yi,t' + Yj,t − Yj,t')] / [Σ j∈S unit, t'∈S time Ai,t' Aj,t Aj,t']

where η = (η1, η2). The numerator terms combine: (i) outcome of unit i at similar time t', (ii) outcome of similar unit j at time t, (iii) minus outcome of unit j at time t'.

The paper also provides a confidence interval for θi,t:


[θ̂ DR i,t,η − z α/2σ̂/√Ji,t,η, θ̂ DR i,t,η + z α/2σ̂/√Ji,t,η]

where Ji,t,η is defined in equation (13) and σ̂2 is a consistent estimate of noise variance.

Assumption 1: Latent factorization θi,t = f(ui, vt)
Assumption 2: Missingness pattern: Ai,t drawn i.i.d. from Bernoulli(p), independent of other randomness
Assumption 3: Bounded noise: E[εi,tU,V] = 0, Var(εi,t) = σ2, εi,t ≤ cε almost surely
Assumption 4: Non-linear factor model: f is twice continuously-differentiable, Lipschitz, with supf ≤ D, sup∇2f ≤ LH, and satisfies convexity conditions on distances (equation 15)

Under Assumptions 2–4, with probability at least pnn, the DR-NN estimate satisfies:


(θ̂ DR i,t,η − θi,t)2 ≤ B DR η + V DR η

where the bias term is:


B DR η = 32L2H [(η'1 + 2eδ/(p√T)) 4/α1/cf,1 + (η'2 + 2eδ/(p√N)) 4/α2/cf,2]

and the variance term is:


V DR η = cσ2/p [log(S t,η2'/δ)/S i,η1' + log(S i,η1'/δ)/S t,η2' + log(S i,η1'S t,η2'/δ)/(p2S i,η1'S t,η2')]

Under Assumptions 2, 3, 5 (bilinear f(u,v) = ⟨u,v⟩ with covariance matrices having positive minimum eigenvalues), the bias simplifies to:


B DR η = (1/(λunit λtime)) [(η'1 + 2eδ/(p√T))2 + (η'2 + 2eδ/(p√N))2]

For T = N and p = Θ(N(−β)):

Example 1 (Discrete factors): With M distinct factor values, error scales as Õ(M/N(1−β) + M2/N(2−3β))
Example 2 (Continuous factors): With β = 0, error scales as Õ(N(−4/(d+4)))

The paper notes: Notably M/N resembles a parametric rate, while N(−4/(d+4)) denotes a non-parametric rate that suffers from the curse of dimensionality. For vanilla NN, the corresponding rates are Õ(N(−1/2) + M/N) and Õ(N(−2/(d+2))), showing DR-NN provides a quadratic and near-quadratic in N improvement in the error.

Under suitable scaling conditions:

(a) Consistency: If B DR η,asymp → 0 and the variance condition (33) holds, then θ̂ DR i,t,η → θi,t in probability.
(b) Asymptotic normality: If conditions (34) and (35) hold, then √Ji,t,η(θ̂ DR i,t,η − θi,t) ⇒ N(0, σ2).

The paper shows that the DR-NN bias is approximately the product of the biases of unit-NN and time-NN, while the variance is approximately the sum of their variances:


B DR η ≈ B unit η1 · B time η2 and V DR η ≈ V unit η1 + V time η2

This leads to two conclusions:

Double robustness: The DR-NN error is at least as good as the better of unit-NN or time-NN, even when one of them fails due to heterogeneity.
Improved performance when both work: When both vanilla estimates provide vanishing error, the DR-NN with optimally tuned hyperparameters provides a significant improvement, with error dominated by the variance sum which is much smaller than either individual variance.

Triply robust NN for tensors: The paper sketches an extension to tensor completion (N × T × M) where the estimator combines three sets of neighbors and provides robustness to heterogeneity in any of the three factor sets.
Alternative derivations: DR-NN can be viewed as (i) a bias-adjusted vanilla NN (using one set of neighbors to estimate bias for the other), and (ii) a minimizer of an orthogonal loss function (first-order approximation of the squared loss).
Connection to SDID: DR-NN can be interpreted as a generalization of the SDID estimator to obtain unit-time level treatment effect guarantees for a wide range of treatment patterns.
The analysis uses a 2×2 data split (time points split into T1, T2; units split into N1, N2) to simplify bias analysis.
The paper assumes MCAR (missing completely at random) missingness, standard in the NN literature.
The bounded noise assumption can be relaxed to sub-Gaussian noise for high-probability results or bounded high-order moments for asymptotic results.
The regularity condition (18) ensures the hyperparameters are chosen so that the effective neighbor sets are non-empty and concentration bounds hold.

Improvements for AI systems

Based on the scientific paper, here are the specific improvements I can make to AI systems and what the improved AI system can do:

Improvement: Replace standard user-user or item-item nearest neighbor collaborative filtering with the doubly robust nearest neighbor (DR-NN) estimator.

What the improved AI system can do:

Provide accurate recommendations even when a user has very few similar users (heterogeneous user base) or when an item has very few similar items (niche items)
Achieve near-quadratic error reduction (e.g., from O(N−1/2) to O(N−1) for discrete latent factors) compared to vanilla nearest neighbor approaches
Maintain performance guarantees even when only one of the two dimensions (users or items) has good neighbors, making it robust to sparse or skewed interaction data

Abstract

We introduce and analyze an improved variant of nearest neighbors (NN) for estimation with missing data in latent factor models. We consider a matrix completion problem with missing data, where the (i, t) -th entry, when observed, is given by its mean f(u i, v t) plus mean-zero noise for an unknown function f and latent factors u i and v t. Prior NN strategies, like unit-unit NN, for estimating the mean f(u i, v t) relies on existence of other rows j with u j about u i. Similarly, time-time NN strategy relies on existence of columns t' with v t' about v t. These strategies provide poor performance respectively when similar rows or similar columns are not available. Our estimate is doubly robust to this deficit in two ways: (1) As long as there exist either good row or good column neighbors, our estimate provides a consistent estimate. (2) Furthermore, if both good row and good column neighbors exist, it provides a (near-)quadratic improvement in the non-asymptotic error and admits a significantly narrower asymptotic confidence interval when compared to both unit-unit or time-time NN.