PaperlayerView exact revision on arXiv

Source-linked AI summary

Demystifying Fixed k-Nearest Neighbor Information Estimators

Weihao Gao, Sewoong Oh, Pramod Viswanath

arXiv:1604.03006v2cs.LGcs.ITstat.ML

TL;DR

Fixed-k KSG mutual-information estimation lacked established consistency and bias-rate guarantees despite widespread use. The paper analyzes KSG and KL nearest-neighbor estimators, proves consistency and convergence results, and introduces BI-KSG through correlation boosting. BI-KSG improves empirical performance especially for small k and moderate N, while parametric entropy rates are achieved for d ≤ 2.

  • Problem

    Fixed-k KSG mutual-information estimation lacked established consistency and bias-rate guarantees despite widespread use.

  • Method

    The paper analyzes fixed-k nearest-neighbor estimators, proving KSG consistency and deriving bias and ℓ2 convergence bounds alongside KL entropy results.

  • Results

    BI-KSG empirically improves on KSG and 3KL, particularly for small k and moderate N, while retaining the same asymptotic theoretical guarantees as KSG.

  • Takeaways & Limitations

    The results establish parametric entropy convergence for d ≤ 2 and nearly tight fixed-k nearest-neighbor error rates in the studied settings.

  • Takeaways & Limitations

    The analysis does not close the upper–lower MSE-rate gap for dimensions d > 2 and excludes some density classes handled by related work.

Abstract

from arXiv · show

Estimating mutual information from i.i.d. samples drawn from an unknown joint density function is a basic statistical problem of broad interest with multitudinous applications. The most popular estimator is one proposed by Kraskov and Stögbauer and Grassberger (KSG) in 2004, and is nonparametric and based on the distances of each sample to its $k^{\rm th}$ nearest neighboring sample, where $k$ is a fixed small integer. Despite its widespread use (part of scientific software packages), theoretical properties of this estimator have been largely unexplored. In this paper we demonstrate that the estimator is consistent and also identify an upper bound on the rate of convergence of the bias as a function of number of samples. We argue that the superior performance benefits of the KSG estimator stems from a curious "correlation boosting" effect and build on this intuition to modify the KSG estimator in novel ways to construct a superior estimator. As a byproduct of our investigations, we obtain nearly tight rates of convergence of the $\ell_2$ error of the well known fixed $k$ nearest neighbor estimator of differential entropy by Kozachenko and Leonenko.

1 Introduction

The paper establishes theoretical guarantees for fixed-k nearest-neighbor estimators of mutual information and differential entropy, addressing longstanding uncertainty about KSG consistency and convergence. It also explains KSG’s empirical advantage through correlation boosting and proposes BI-KSG, while reporting parametric entropy rates in low dimensions.

  • KSG estimator: The KSG estimator is consistent, with bias upper bound eO(N^-1/(d_x+d_y)) and variance eO(1/N).Its ℓ2 error is eO(1/√N + N^-1/(d_x+d_y)).
  • Correlation boosting and BI-KSG: BI-KSG amplifies KSG’s correlation-boosting effect and improves empirical performance, especially for small k and moderate N.Its asymptotic theoretical guarantees are the same as KSG’s.
  • Correlation boosting and BI-KSG: The paper identifies correlation boosting as the mechanism behind KSG’s improved performance over 3KL and uses it to motivate a superior estimator.The motivation is grounded in KSG’s widespread use despite previously limited theoretical understanding.
  • KL entropy estimator: Fixed-k nearest-neighbor distances provide local density information used to construct the KL differential-entropy estimator and related mutual-information estimators.The KL estimator includes a digamma correction that is crucial for debiasing when k is a small constant.
  • KL entropy estimator: For differential entropy, the KL estimator achieves the parametric ℓ2 error rate eO(1/√N) when d ≤ 2.The result applies to a class of bounded-support densities, including uniform and truncated Gaussian distributions.
  • KSG estimator: KSG modifies the 3KL mutual-information estimator through sample-dependent marginal nearest-neighbor distances that exactly cancel distance terms.This construction motivated the paper’s analysis because KSG’s empirical superiority lacked a theoretical explanation.

2 Convergence Rate of KL Entropy Estimator

The paper derives upper bounds for the KL estimator’s bias, variance, and MSE under smoothness, support, and boundary assumptions, and compares them with minimax lower bounds. The bounds match in low dimensions, while a gap remains for d > 2 and empirical rates are close to the theoretical upper bounds.

  • Upper bounds: The analysis generalizes prior convergence guarantees from k = 1 and d = 1 to arbitrary fixed k and general dimensions under technical assumptions.The assumptions include smoothness, tail control, and a finite boundary surface measure.
  • Assumptions and boundary effects: Finite boundary surface area limits boundary bias to approximately O~(N^-1/d), which dominates the approximately O~(N^-2/d) interior bias.The boundary restriction is needed because k-NN distances become larger near boundaries, increasing density-estimation bias.
  • Upper bounds: The KL estimator’s MSE upper bound combines separate bias and variance bounds for fixed k and dimension d under Assumption 1.The variance bound is dimension-independent, while the resulting MSE bound includes dimension-dependent bias terms and logarithmic factors.
  • Comparing the bounds: Theoretical upper and lower MSE bounds match for d ≤ 2 at the parametric rate, but an exponent gap remains for d > 2.Closing the higher-dimensional gap is identified as a future research direction.
  • Empirical evidence: For uniform distributions, simulated KL-estimator rates are fairly close to the theoretical upper bounds, while empirical k-NN performance can exceed KDE-based methods at modest sample sizes.The latter comparison is illustrated for independent Beta(2,2) variables in six dimensions over 500 trials.

3 KSG Estimator: Consistency and Convergence Rate

The paper establishes consistency for KSG and BI-KSG under regularity assumptions and derives convergence and MSE bounds for fixed-k mutual-information estimation. In the one-dimensional balanced case, the KSG, BI-KSG, and 3KL estimators achieve the parametric MSE rate up to polylogarithmic factors.

  • 3.1 Consistency: The KSG and BI-KSG estimators are consistent in probability under the stated assumptions.The assumptions impose integrability, bounded conditional densities, smoothness, and related regularity conditions.
  • 3.2 Convergence rate: The formal convergence results use truncated k-nearest-neighbor distances and define corresponding truncated KSG and BI-KSG estimators.The truncation threshold depends on the dimensions dx and dy and a parameter δ.
  • 3.2 Convergence rate: The truncated KSG and BI-KSG estimators have MSE upper bounds obtained by combining their bias and variance bounds.The bounds apply under Assumption 3 for finite k = O(1) and d = O(1).
  • 3.2 Convergence rate: The upper bounds for 3KL, KSG, and BI-KSG have the same MSE scaling.This comparison concerns the rates represented by the respective upper bounds.
  • 3.2 Convergence rate: 1/N MSE convergence holds for KSG, BI-KSG, and 3KL when dx = dy = 1, up to a poly-logarithmic factor.The paper identifies this as the parametric rate, which cannot be improved.

4 Correlation Boosting

The section explains KSG’s advantage through positively correlated joint and marginal entropy-estimation biases, then uses this mechanism to motivate BI-KSG’s distance and count transformations.

  • Correlation Boosting: Positive correlation among joint and marginal entropy-estimation biases reduces KSG’s mutual-information bias, with stronger correlation than 3KL and sometimes BI-KSG.The authors hypothesize that this correlation boosting explains KSG’s lower mean-square error than 3KL.
  • New Estimator: BI-KSG replaces KSG’s ℓ∞ distances with ℓ2 distances and replaces digamma terms with logarithms of marginal neighbor counts.The logarithmic replacement removes the O(1/nx,i,2 + 1/ny,i,2) bias identified for the corresponding entropy construction.
  • Correlation Boosting: KSG’s ℓ∞ neighborhood can make the shared radius an exact marginal neighbor distance for one variable but not the other, creating bias of order 1/nx,i,∞ or 1/ny,i,∞.Which marginal is affected depends on whether the joint nearest neighbor lies on the X- or Y-boundary of the hypercube.
  • New Estimator: BI-KSG is consistent with bias eO(N^(-1/(dx+dy))) and variance eO(1/N), giving ℓ2 error eO(1/√N + N^(-1/(dx+dy))).Its asymptotic theoretical guarantees match those established for KSG.
  • Empirical Evaluation: For k = 1 and correlated Gaussian variables, all estimators approach the ground truth, while BI-KSG has the best sample complexity at moderate N.The reported empirical gains are strongest for moderate sample sizes, when marginal neighbor counts remain small.

5 Multivariate Mutual Information

The section extends correlation-boosting ideas from mutual information to multivariate mutual information and balanced set-function measures, producing estimators that outperform KL-based alternatives empirically.

  • Multivariate Mutual Information: Multivariate mutual information can be estimated either with an (L + 1)-KL construction or with KSG- and BI-KSG-style correlation-boosting estimators.The correlation-boosting approach generalizes the basic KSG constructions to sums and differences of entropy terms.
  • Multivariate Mutual Information: Correlation-boosting estimators perform significantly better than the simpler (L + 1)-KL estimator for three jointly Gaussian variables across N = 100–3000.The comparison uses the MSE-versus-sample-size experiment shown in Figure 6.
  • Balanced Measures: Balanced set-function measures assign coefficients so each variable’s total positive and negative contributions cancel, enabling correlation-boosting estimation from shared nearest-neighbor distances.The paper illustrates this procedure for a four-variable causal-strength-related measure.
  • Balanced Measures: The four-variable example estimates joint entropy with a joint-space k-NN radius and uses that radius to count neighbors in each relevant subset.This construction preserves the balanced cancellation structure across the entropy terms.

6 Related Work

The paper situates fixed k-NN information estimation within continuous nonparametric estimation, emphasizing geometry, dimensionality, and smoothness. It contrasts existing approaches while identifying limits of standard smoothness-based refinements for fixed-k estimators.

  • Continuous information estimation depends on the interplay between Euclidean geometry and the dimensionality of the domain.
  • Existing alternatives include KDE with resubstitution or data splitting, spacings, Edgeworth expansions, convex optimization, and k-NN estimators with growing or fixed k.
  • KDE-based differential entropy estimation can achieve the parametric rate when the dimension is no more than 6.
  • The analysis studies bounded-Hessian and Hölder smoothness classes, but standard Hölder modeling does not yield refined rates for fixed-k KL and KSG estimators.
  • The analysis covers bounded-Hessian settings while leaving open how to unify cases involving nonsmooth boundaries and densities approaching zero.

7 Proof of Theorem 3

The proof of Theorem 3 analyzes the fixed-k KL estimator by representing its nearest-neighbor statistic through distributional approximations and decomposing the bias into separately bounded terms. The resulting bounds establish the theorem’s convergence behavior under the stated assumptions.

  • The proof establishes a generalization of the KL entropy estimator’s convergence rate to arbitrary dimensions and arbitrary fixed k.The argument also applies to any ℓp distance for 1 ≤ p ≤ ∞.
  • The nearest-neighbor statistic is reparameterized through a truncated radius transformation, allowing its expectation to be expressed using a transformed variable.
  • The bias is decomposed into three terms, each controlled separately through bounds on truncation, density approximation, and distributional differences.
  • The final bias bound follows by combining the bounds for the three decomposed terms under the paper’s assumptions.
  • The proof approximates local probability mass using density times ball volume, with bounded Hessians controlling the approximation error over relevant regions.

8 Proof of Theorem 4

The proof of Theorem 4 controls the KL estimator’s variance by separating individual variance terms from pairwise covariance terms. It bounds dependence through local-neighborhood overlap and disjoint-ball approximations.

  • The variance is decomposed into a single-sample variance term and a covariance term between two nearest-neighbor statistics.
  • Combining the near- and far-location bounds yields the desired covariance control for the variance analysis.
  • Conditional variances are bounded using the transformed nearest-neighbor densities and logarithmic moment controls.
  • When sample locations are close, covariance is controlled with Cauchy–Schwarz and bounds on the conditional variances.
  • When locations are sufficiently separated, their nearest-neighbor balls are disjoint, enabling comparison of the joint order-statistic density with the product of marginal densities.

9 Proof of Theorem 5 on the minimax lower bound

The proof derives minimax lower bounds for differential entropy estimation using Le Cam’s method. It combines simple Gaussian alternatives with smooth localized perturbations to obtain dimension-dependent difficulty.

  • Le Cam’s method is used to establish an Ω(1/N) minimax lower bound from two Gaussian distributions with bounded Hessian norms.
  • The lower-bound construction uses a uniform density on [0,1]^d and localized perturbations indexed by sign vectors.
  • The perturbation family is chosen within the smoothness class while keeping the averaged N-sample distributions close in total variation.
  • The resulting minimax lower bound is N^-16/(d+8) when the construction uses γ = 2.

10 Proof of Theorem 6 on the consistency of KSG estimator

The proof establishes consistency of the KSG and BI-KSG mutual-information estimators by combining convergence of marginal and joint entropy estimates. Under the stated condition k > d_y/d_x, all error terms vanish as N grows, yielding convergence in probability.

  • Consistency of entropy estimates: The joint entropy estimates for KSG and BI-KSG converge in probability, while the marginal estimates converge by the technical lemma and its X/Y symmetry.The argument uses known convergence for the joint estimates and proves the marginal case directly.
  • Error decomposition: The proof decomposes the remaining error into empirical-mean and density-estimation components, then controls each through concentration and local-density bounds.The density-estimation component is analyzed using the k-nearest-neighbor radius and separate bounds for three terms.
  • Error decomposition: When k > d_y/d_x, each of the three error terms tends to zero as N approaches infinity.This condition is used to complete the convergence argument for the marginal estimator.
  • Consistency of entropy estimates: The marginal entropy estimate converges to the true entropy in probability after combining sampling and density-estimation errors.The proof bounds three error contributions and shows their upper bounds vanish as N approaches infinity.

11 Proof of Theorem 7 on the bias of KSG estimator

The proof bounds the bias of the KSG estimator by controlling large-neighborhood events, local entropy-estimation errors, and marginal entropy bias. Under the stated regularity assumptions and k ≥ d_x/d_y, these bounds establish the theorem’s convergence-rate claim.

  • Bias decomposition: The bias analysis combines bounds for joint and marginal entropy estimators with control of the k-nearest-neighbor radius exceeding a threshold.The proof splits the support and radius events, then applies the corresponding entropy and tail bounds.
  • Conclusion: Combining the joint-entropy result, marginal bias lemma, and tail bound yields the desired convergence-rate statement for the KSG estimator.The same argument applies to the Y marginal after exchanging X and Y.
  • Local density control: The proof controls local density approximation errors using smoothness of the density and a decomposition into regular and high-curvature regions.The neighborhood probability is compared with the local density times ball volume, while the exceptional region is bounded separately.
  • Small-radius control: For k ≥ d_x/d_y, the small-radius event satisfies P(ρ_k,1 < r_N) ≤ C(log N)^(2k/d_x)/N.Here r_N = (log N)^2N^(-1/d_x), and the bound controls the contribution from unusually small neighbor distances.
  • Small-radius control: The corresponding logarithmic-radius integral is bounded by C(log N)^(2k/d_x+1)/N.This bound is inserted into the decomposition of the marginal entropy bias.

12 Proof of Theorem 8 on the variance of KSG estimator

The variance proof reduces KSG variance to marginal entropy variances and covariances, then analyzes covariance through three geometric configurations of neighbor regions. Bounds for each configuration are combined to complete the theorem.

  • Variance decomposition: The proof reduces mutual-information variance to marginal entropy variances and covariances using the KSG decomposition and Cauchy–Schwarz.The marginal variance is represented through individual-estimator variance and pairwise covariance terms.
  • Geometric covariance cases: The covariance analysis distinguishes disjoint strips, overlapping strips with disjoint joint balls, and intersecting joint balls.These three cases determine whether multinomial calculations, conditional independence, or Cauchy–Schwarz bounds apply.
  • Disjoint strips: For disjoint strips, the count variables are multinomial and the covariance is bounded by 3/N for sufficiently large N.The multivariate delta method is applied to logarithmic count functions.
  • Overlapping strips: For overlapping strips with disjoint joint balls, the covariance is bounded by 4(p_1+p_2)/N.The shared region R2 enters through p_2, while p_1+p_2 is the probability in the first strip outside its joint ball.

13 Proof of Theorem 9

The proof characterizes marginal neighbor counts conditional on a joint k-nearest-neighbor radius. It shows that the remaining counts are sums of i.i.d. Bernoulli variables, hence binomially distributed.

  • Conditional count distribution: The proof partitions sample indices into k−1 closer points, the kth neighbor, and remaining farther points before conditioning on the kth-neighbor event.This partition identifies which observations contribute deterministically and which contribute Bernoulli indicators.
  • Bernoulli representation: Given the conditioning event, the remaining X-distance indicators are i.i.d. Bernoulli variables with a common probability determined by z and r.Independence follows after conditioning on the joint-distance constraints for the remaining samples.
  • Conditional count distribution: Conditional on Z_i = z and ρ_k,i = r, the marginal count n_x,i − k has a binomial distribution.The Bernoulli trials correspond to remaining samples whose X-distance from x is smaller than r.
Loading 1604.03006v2…