[Paper Review] Sublinear Partition Estimation

Pushpendre Rastogi, Benjamin Van Durme|arXiv (Cornell University)|Aug 7, 2015

Advanced Image and Video Retrieval Techniques17 references3 citations

TL;DR

This paper proposes three sublinear methods to estimate the partition function in neural network classifiers, using approximate nearest neighbor search and kernel feature maps to reduce computation time. The MIMPS method achieves high accuracy with only 100 head and 100 tail samples, outperforming the common NCE heuristic of assuming Z ≈ 1.

ABSTRACT

The output scores of a neural network classifier are converted to probabilities via normalizing over the scores of all competing categories. Computing this partition function, $Z$, is then linear in the number of categories, which is problematic as real-world problem sets continue to grow in categorical types, such as in visual object recognition or discriminative language modeling. We propose three approaches for sublinear estimation of the partition function, based on approximate nearest neighbor search and kernel feature maps and compare the performance of the proposed approaches empirically.

Motivation & Objective

• To address the computational bottleneck of computing the partition function Z(q) = ∑exp(ui) in large-scale neural network classifiers with tens of thousands of output classes.
• To reduce the O(N) complexity of brute-force partition function computation to sublinear time, enabling faster inference in large-vocabulary tasks like language modeling and image recognition.
• To develop practical, empirically effective estimators that avoid the need for external models or training-time constraints on Z.
• To evaluate the performance of these estimators under controlled and real-world conditions, particularly in language modeling.

Proposed method

• Proposes MIMPS (MIPS-based Importance Sampling), a method that uses approximate nearest neighbor search to identify the k highest-scoring classes and estimates Z(q) using importance sampling with these samples.
• Introduces MINCE (MIPS-based NCE), which combines nearest neighbor retrieval with Noise Contrastive Estimation to estimate Z(q) using a small set of high-scoring and low-scoring classes.
• Employs kernel feature maps to transform the score space into a higher-dimensional space where inner products better approximate exponential scores, enabling efficient estimation.
• Uses a two-stage sampling strategy: 'head' samples from the top-k nearest neighbors and 'tail' samples from the remainder to balance variance and bias in the estimator.
• Implements the MIPS-based approach using the FLANN library with a modified K-Means Tree algorithm for efficient retrieval.
• Employs controlled experiments with synthetic data to isolate the impact of retrieval error and sample size on estimation accuracy.

Experimental results

Research questions

• RQ1 Can approximate nearest neighbor search be effectively leveraged to estimate the partition function in sublinear time?
• RQ2 How does retrieval error in nearest neighbor algorithms affect the accuracy of partition function estimation?
• RQ3 Does the MIMPS estimator outperform the standard NCE heuristic of assuming Z ≈ 1 in real-world language modeling tasks?
• RQ4 How do the number of head and tail samples influence the trade-off between estimation accuracy and computational speed?
• RQ5 Can kernel feature maps improve the quality of partition function estimation when combined with nearest neighbor sampling?

Key findings

• With 100 head and 100 tail samples, the MIMPS estimator achieved an absolute error of 198.5 on the Penn Treebank test set, significantly outperforming the NCE heuristic’s 352 error.
• The MIMPS method was better than the NCE heuristic in 70.5% of test contexts when using 100 head and 100 tail samples, demonstrating consistent improvement.
• The method achieved a 10x speedup over brute-force computation while maintaining high estimation accuracy.
• Controlled experiments showed that estimation error increases with retrieval error, highlighting the importance of accurate nearest neighbor indexing.
• The MINCE estimator performed poorly in practice, failing to achieve reliable estimates despite theoretical promise.
• The results suggest that the performance of sublinear partition estimation critically depends on the quality of the underlying nearest neighbor indexing mechanism.