Scalable Random Wavelet Features: Efficient Non-Stationary Kernel Approximation with Convergence Guarantees

AI summary

Overview

Research area: Machine learning — scalable kernel methods, Gaussian process approximation, and non-stationary function modeling.

Technical level: Intermediate. The paper assumes familiarity with Gaussian processes, kernel functions, and random feature approximations, though the core idea (replacing sine waves with small localized wavelets) is conceptually accessible.

Scope: The paper introduces Random Wavelet Features (RWF), a random-feature framework built from wavelet atoms rather than Fourier bases, and evaluates it on synthetic, speech, and regression benchmarks against stationary and non-stationary GP baselines.

What This Paper Is About

Most scalable kernel methods, such as Random Fourier Features (RFF), rely on Bochner's theorem, which only applies to stationary kernels — kernels whose behavior is uniform across the input space. This means they systematically misspecify data whose statistical properties change from region to region, like terrain that shifts abruptly across a map or speech signals with sudden consonant bursts. The authors' goal is to build a random feature map that is as cheap to train as RFF but can represent non-stationary structure, by sampling randomly scaled and shifted wavelets instead of globally supported sinusoids.

Key Contributions

1. A non-stationary kernel construction from wavelets. The kernel is defined as an integral over products of wavelet atoms ψ_{s,t}(x)ψ_{s,t}(y) with respect to a non-negative measure over scales s and shifts t, and the authors prove this integral yields a positive definite kernel (Theorem 4.1).

2. Monte Carlo feature map with theoretical guarantees. Sampling D i.i.d. (s_i, t_i) pairs gives an explicit feature map z(x) = (1/√D)[ψ_{s_1,t_1}(x), …, ψ_{s_D,t_D}(x)]^T whose inner product is an unbiased estimator of the kernel (Lemma 4.1), with a variance bound of B²/D (Lemma 4.2) and a uniform convergence bound over compact sets (Theorem 4.2), including an explicit sample complexity expression for D.

3. Linear-time training complexity. Constructing D features over N inputs costs O(NDd); forming Z^T Z costs O(ND²); the resulting D × D system solves in O(D³); so for N ≫ D, training is dominated by O(ND²), and prediction costs O(D²) per test point. This is contrasted with exact GPs at O(N³) and SVGP at O(NM²) per optimization step.

4. Empirical evaluation across synthetic, speech, and benchmark data. RWF is compared against Exact GP, SVGP, RFF-GP, Deep GP, Deep-RF GP (DRF), Spectral Mixture kernels, Inter-domain Deep GP (IDD), and Adaptive RKHS Fourier Feature GPs (A-RKHS) on a multi-step function, TIMIT speech regression, seven UCI datasets, and the Protein dataset.

Main Findings

• Multi-step synthetic function (Table 1, five runs): RWF achieved the lowest RMSE (0.071 ± 0.011), CRPS (0.112 ± 0.010), and NLL (-1.879 ± 0.061), and the shortest reported time (9) among the compared methods. The next best RMSE was A-RKHS at 0.095 ± 0.045.

• TIMIT speech regression (Table 2, five runs): RWF reported the lowest RMSE (0.42 ± 0.003), CRPS (0.39 ± 0.006), and NLL (0.56 ± 0.04) with the shortest training time (90), versus A-RKHS at 0.48 ± 0.003 RMSE and 141 time units, and RFF at 2.13 ± 0.004 RMSE with 110 time units.

• UCI regression (Table 3): RWF reported the lowest RMSE on five of the seven datasets (ENERGY 0.42 ± 0.02, AIRFOIL 3.20 ± 0.08, STOCK 0.84 ± 0.03, KIN8NM 0.09 ± 0.01, NAVAL 0.02 ± 0.001) and competitive results on the remaining two. It also reported the lowest training time on all seven (9, 8, 9.6, 10, 12, 9, and 24 minutes).

• Protein dataset (Table 4, approximately 45K examples, 9 input features): RWF reported the best RMSE (4.25 ± 0.02), CRPS (3.65 ± 0.02), and NLL (2.71 ± 0.03) and the minimum training time (90 minutes). The nearest RMSE was A-RKHS at 4.32 ± 0.01.

• Theoretical scaling: The sample complexity result shows the required number of features D grows logarithmically in the accuracy and confidence terms through the log(4 diam(M) L_z / ε) and log(2/δ) factors, while the leading dependence on accuracy is 8B²/ε².

• Qualitative behavior: The paper reports that RWF-GP captures sharp discontinuities in the multi-step function while maintaining calibrated uncertainty, whereas Exact GP, SVGP, and RFF-GP either over-smooth or miscalibrate uncertainty. Figure 2 reports RWF as most efficient in time and memory as training samples grow, with SVGP and Deep GP incurring higher cost.

• Not reported: Exact numerical values behind Figures 1 and 2 (time and memory curves versus number of training samples) are not given in the truncated text, and the appendix converged ablation on feature size (Section C.1) is referenced but not included in the provided content.

Methodology in Plain English

The authors start from a mother wavelet ψ — a function with zero mean and unit L² norm — and generate a family of "wave packets" by scaling it by s and shifting it by t: ψ_{s,t}(x) = s^{-d/2} ψ((x − t)/s). Small scales produce narrow, high-frequency atoms that track sharp local events; large scales produce broad, low-frequency atoms that track smooth trends.

They then define a kernel as an integral of products of these atoms over all scales and shifts, weighted by a non-negative density p(s, t). Because the integrand is a product of scalar features, the resulting kernel is automatically positive definite. A typical factorized choice is p(s, t) = p_s(s) p_t(t), with a log-uniform distribution over scales to span multiple resolutions and a uniform distribution over shifts covering the data's convex hull.

Since that integral cannot be computed in closed form, the method draws D independent (s_i, t_i) samples and stacks the resulting atoms into a D-dimensional feature vector. The kernel approximation is just the dot product of two such vectors, which turns the GP into a Bayesian linear regression over wavelet features. Training then reduces to computing the standard Gaussian posterior over the weights via a closed-form D × D solve, and the same two equations used in RFF-GP give predictive mean and variance. Hyperparameters (noise variance and the parameters of the sampling distribution) are optimized by maximizing the marginal likelihood of that linear model.

The theory follows by bounding the feature magnitude with B = sup|ψ_{s,t}(x)| and assuming Lipschitz continuity of the kernel and feature map, then applying a covering-number argument over a compact set of diameter diam(M).

Why This Matters

Impact on research. The work supplies a way to get random-feature scalability — linear-in-data training — without the stationarity assumption that Bochner's theorem imposes. It also fills a gap in the literature the authors identify: prior wavelet-based kernel work relied on fixed or predefined wavelet dictionaries and lacked Monte Carlo approximation or guarantees like unbiasedness and uniform convergence. The explicit sample complexity bound gives practitioners a principled way to choose the number of features.

Real-world applications:

• Geospatial modeling, where terrain and environmental statistics vary sharply across regions, a case the authors name directly in the introduction.
• Speech and audio analysis, where signals contain localized consonant bursts alongside slowly varying regions, tested here on a TIMIT-derived regression task.
• Large-scale scientific regression, such as the Protein dataset (approximately 45K examples, 9 features) drawn from a noisy biological domain.
• UCI-style tabular regression across energy, materials, financial, motion, kinematics, and naval datasets, spanning sample sizes from 1k to 11k.

Industry relevance. Anyone already using RFF-GP as a fast surrogate will find the RWF machinery a drop-in variant with the same posterior update equations and lower reported wall-clock times across every table. The reported single closed-form posterior solve, versus the iterative variational updates of SVGP, is the main practical selling point for repeated retraining or deployment on large datasets.

Future Directions

• Adaptive wavelet sampling: the conclusion names adaptive wavelet sampling as a direction for further expanding RWF's reach, suggesting the current sampling distribution p(s, t) is fixed rather than input-adaptive.
• Integration with deep kernel architectures: the authors propose combining RWF with deep kernel models, which would stack wavelet features rather than using them as a single shallow layer.
• Choice of wavelet and sampling distribution: the paper notes that the mother wavelet (for example Morlet for time-frequency analysis or Daubechies for sharp transitions) and the scale/shift distributions are practical design decisions that shape performance, leaving a systematic study open.
• Understanding when localization helps most: the UCI results show RWF leading on five of seven datasets but not all, so the question of which data regimes benefit most from wavelet localization versus other non-stationary constructions remains open.

Target Audience

Researchers and practitioners in probabilistic machine learning who work with Gaussian processes or kernel methods at scale, particularly those who have hit the limits of RFF-GP on data with region-dependent behavior. The paper is also relevant to applied scientists in geospatial modeling, speech processing, and computational biology who need a surrogate model that is both expressive and cheap to retrain. Readers without background in GP regression or kernel theory will need to consult the preliminaries section and appendix, as the theory portion assumes comfort with positive definite kernels, covering arguments, and Lipschitz bounds.

Authors’ abstract

Modeling non-stationary processes, where statistical properties vary across the input domain, is a critical challenge in machine learning; yet most scalable methods rely on a simplifying assumption of stationarity. This forces a difficult trade-off: use expressive but computationally demanding models like Deep Gaussian Processes, or scalable but limited methods like Random Fourier Features (RFF). We close this gap by introducing Random Wavelet Features (RWF), a framework that constructs scalable, non-stationary kernel approximations by sampling from wavelet families. By harnessing the inherent localization and multi-resolution structure of wavelets, RWF generates an explicit feature map that captures complex, input-dependent patterns. Our framework provides a principled way to generalize RFF to the non-stationary setting and comes with a comprehensive theoretical analysis, including positive definiteness, unbiasedness, and uniform convergence guarantees. We demonstrate empirically on a range of challenging synthetic and real-world datasets that RWF outperforms stationary random features and offers a compelling accuracy-efficiency trade-off against more complex models, unlocking scalable and expressive kernel methods for a broad class of real-world non-stationary problems.