[Paper Review] On Feature Collapse and Deep Kernel Learning for Single Forward Pass Uncertainty

Joost van Amersfoort, Lewis Smith+3 more|arXiv (Cornell University)|Feb 22, 2021

Gaussian Processes and Bayesian InferenceComputer Science48 references56 citations

TL;DR

The paper identifies feature collapse in Deep Kernel Learning (DKL) for single forward-pass uncertainty and proposes Deterministic Uncertainty Estimation (DUE) with bi-Lipschitz constrained feature extractors and inducing-point GPs to achieve superior uncertainty estimates while remaining fast.

ABSTRACT

Inducing point Gaussian process approximations are often considered a gold standard in uncertainty estimation since they retain many of the properties of the exact GP and scale to large datasets. A major drawback is that they have difficulty scaling to high dimensional inputs. Deep Kernel Learning (DKL) promises a solution: a deep feature extractor transforms the inputs over which an inducing point Gaussian process is defined. However, DKL has been shown to provide unreliable uncertainty estimates in practice. We study why, and show that with no constraints, the DKL objective pushes "far-away" data points to be mapped to the same features as those of training-set points. With this insight we propose to constrain DKL's feature extractor to approximately preserve distances through a bi-Lipschitz constraint, resulting in a feature space favorable to DKL. We obtain a model, DUE, which demonstrates uncertainty quality outperforming previous DKL and other single forward pass uncertainty methods, while maintaining the speed and accuracy of standard neural networks.

Motivation & Objective

• Motivate the need for reliable uncertainty estimates in fast, single forward-pass models.
• Diagnose why standard DKL-based uncertainty can be unreliable due to feature collapse.
• Propose a constrained DKL approach (DUE) with bi-Lipschitz feature extractors to improve uncertainty quality.
• Demonstrate that DUE achieves competitive or superior uncertainty performance while maintaining neural-network speed and accuracy.

Proposed method

• Analyze feature collapse in DKL when the deep feature extractor is unconstrained.
• Impose bi-Lipschitz constraints on the feature extractor via residual connections and spectral normalization to enforce sensitivity and smoothness.
• Place a deep GP (inducing points) on top of the constrained feature extractor for true non-parametric uncertainty estimation.
• Train end-to-end from scratch with practical simplifications (no pre-training, few inducing points).
• Compare against SNGP and other single forward-pass methods using AUROC, accuracy, and predictive uncertainty metrics.
• Provide practical training steps (Algorithm 1) including inducing point initialization via K-means and spectral normalization adjustments.

Experimental results

Research questions

• RQ1 How does feature collapse occur in DKL and how does it affect uncertainty estimates?
• RQ2 Can bi-Lipschitz constraints on the feature extractor mitigate feature collapse and improve uncertainty quality in DKL?
• RQ3 Does a DKL-based single forward-pass model with inducing points (DUE) outperform existing methods on standard uncertainty benchmarks and regression tasks?
• RQ4 Is DUE practical to train from scratch and competitive in speed and accuracy with standard neural networks?
• RQ5 How does DUE perform on CIFAR-10 vs SVHN discrimination and in causal/medical uncertainty benchmarks?

Key findings

• Feature collapse in unconstrained DKL leads to high confidence on out-of-distribution data, degrading uncertainty estimates.
• Bi-Lipschitz constrained feature extractors (residual connections with spectral normalization) mitigate feature collapse and improve uncertainty behavior.
• DUE, using an inducing-point GP on top of a constrained feature extractor, achieves strong uncertainty performance comparable to or better than competing single forward-pass methods on CIFAR-10 vs SVHN.
• DUE trains from scratch with minimal overhead (no pre-training) and small inducing point counts (e.g., 10) for CIFAR-10, with runtime close to a standard softmax model.
• DUE outperforms alternative single forward-pass methods on a CIFAR-10 vs SVHN uncertainty task and a regression benchmark for personalized medicine, while remaining significantly faster than ensembles.
• The inducing-point GP approach preserves non-parametric GP properties and yields uncertainty outside the training support similar to full GPs, unlike RFF-based methods.