Improving TensorSketch Using Complex Random Variables

Improving TensorSketch Using Complex Random Variables

Summary

This paper introduces a novel variant of the TensorSketch algorithm, termed Complex-to-Real (CtR) TensorSketch, which leverages complex random variables to achieve improved variance bounds for polynomial kernel approximation while maintaining the input-sparsity running time of the original TensorSketch method.

Problem Context

Polynomial kernels, defined as k(x, y) = ⟨x, y⟩ p for vectors x, y ∈ R d, are widely used in machine learning to model higher-order non-linear interactions. These kernels are equivalent to mapping x to its p-fold Kronecker product x⊗p ∈ R(d p), where the dimension grows exponentially with p, making direct computation infeasible. Two primary sketching approaches exist for approximating these kernels:

1. JL-type projections (Kar and Karnick [2012]): Dense random projections with computational cost O(pDd), where D is the sketch dimension.

2. TensorSketch (Pham and Pagh [2013]): Combines CountSketch with FFT-based convolution, achieving input-sparsity time O(p(nnz(x) + D log D)).

Both approaches have variance that grows exponentially with the polynomial degree p, scaling as 3 p/D. Recent work by Wacker et al. [2023] showed that using complex-valued distributions reduces this dependence to 2 p/D for the JL-type approach, but their method relies on dense projections and does not extend to TensorSketch.

Key Contribution

The paper proposes a Complex-to-Real (CtR) variant of TensorSketch that:

• Uses random functions drawn independently and uniformly from the fourth roots of unity 1, ω, ω2, ω3 where ω = e(2πi/4) = i

• Combines these with FFT-based tensor sketching

• Produces real-valued embeddings through a complex-to-real conversion

• Achieves the same variance bound as Wacker et al. [2023] (2 p/D dependence) while retaining input-sparsity running time

Technical Details

The method is defined in Definition 8, which constructs p independent Complex CountSketch mappings. Each mapping uses:

• A universal hash function h r: [d] → [D/2] assigning coordinates uniformly to buckets

• A random function s r: [d] → 1, ω, ω2, ω3 drawn uniformly from fourth roots of unity

The sketch is computed via FFT-based convolution:

Φ C(x⊗p):= Cx⊗p = FFT−1(∏ r=1 p FFT(C r x)) ∈ C(D/2)

The CtR sketch is then formed by concatenating real and imaginary parts:

Φ CtR(x⊗p):= [Re Φ C(x⊗p) 1,..., Re Φ C(x⊗p) D/2, Im Φ C(x⊗p) 1,..., Im Φ C(x⊗p) D/2] ∈ R D

Theoretical Guarantees (Theorem 2)

The paper proves that the CtR TensorSketch estimator satisfies:

• Unbiasedness: E[k̂ CtR(x, y)] = ⟨x⊗p, y⊗p⟩ = ⟨x, y⟩ p

• Variance bound: Var(k̂ CtR(x, y)) ≤ ((2(p+1) - 2)/D) · ∥x∥2(2p) · ∥y∥2(2p)

• Time complexity: O(p(nnz(x) + D log D)) for sketching

This improves upon the classical TensorSketch variance bound of (3 p - 1)/D · ∥x∥2(2p) · ∥y∥2(2p), reducing the exponential dependence from 3 p to 2 p.

Key Technical Insights

The paper demonstrates that the variance reduction is non-trivial. The authors prove (Theorem 5 in the Appendix) that using complex random variables in standard CountSketch alone provides no variance reduction over its real-valued counterpart. The improvement in TensorSketch stems from the vanishing of certain cross terms in the variance analysis, which does not occur in the real-valued setting.

The proof relies on Lemma 4, which provides exact moment evaluations for products of independent random functions drawn from fourth roots of unity. This lemma shows that:

• E[Z2] = (⟨x, y⟩2 + ∥x∥22∥y∥22 - Σi xi2yi2) p

• E[Z2] = (2⟨x, y⟩2 - Σi xi2yi2) p

Comparison with Baselines

The paper provides a comprehensive comparison (Table 1) showing:

Method Variance Sketching Time

CtR TensorSketch ((2(p+1)-2)/D)∥x∥2(2p)∥y∥2(2p) O(p(nnz(x)+nnz(y)+D log D))

Real TensorSketch ((3 p-1)/D)∥x∥2(2p)∥y∥2(2p) O(p(nnz(x)+nnz(y)+D log D))

JL CtR Rademacher ((2(p+1)-2)/D)∥x∥2(2p)∥y∥2(2p) O(pDd)

JL CtR Gaussian ((2(p+1)-2)/D)∥x∥2(2p)∥y∥2(2p) O(pDd)

Experimental Results

The paper validates the theoretical findings through experiments on:

1. Synthetic data: n = 3000 random vectors in R2 with Gaussian distribution, polynomial degrees p ∈ 10, 15, 20, sketch dimensions D ∈ 100, 300, 500

2. Real-world datasets: MAGIC Gamma Telescope (d = 10) and COD-RNA (d = 8), with n = 3000 points, degrees p ∈ 3, 5, 7, dimensions D ∈ d, 3d, 5d

Key findings:

• KL divergence: CtR TensorSketch consistently achieves the lowest KL divergence across all degrees and sketch dimensions, indicating more accurate kernel approximation

• Runtime: CtR TensorSketch retains input-sparsity cost, while JL-type methods incur higher dense projection costs

• Variance analysis: Numerical tables confirm CtR TensorSketch achieves lower variance than all baselines across degrees p ∈ 15, 20, 25, 30

• Downstream tasks: In linear SVM classification experiments, CtR TensorSketch achieves the highest accuracy among all sketching baselines while requiring the shortest total execution time

Implications and Applications

The improved variance bounds make CtR TensorSketch a promising alternative for applications including:

• Efficient training of linear SVMs

• Scalable deep learning

• Theoretical analysis of over-parameterized neural networks

• Compact bilinear pooling for fine-grained visual recognition

• Multimodal fusion architectures

• Natural language processing

• Recommender systems

• Genomics

Future Directions

The paper identifies two open research directions:

1. Whether employing higher-order roots of unity can yield stronger higher-moment concentration, leading to tighter (ε, δ)-approximation guarantees

2. Whether the advantages of CtR constructions can be extended to other kernel families and randomized feature maps beyond polynomial kernels

Improvements to AI Systems:

1. Faster and More Accurate Kernel-Based Learning: Integrate CtR TensorSketch into kernel approximation pipelines for SVMs, Gaussian processes, and kernel ridge regression. The improved variance bound (2 p vs 3 p) means AI systems can achieve the same approximation accuracy with smaller sketch dimensions, reducing memory and computational overhead by up to 33% for high-degree polynomial kernels while maintaining input-sparsity time complexity.

2. Scalable Higher-Order Feature Interactions: Enhance feature engineering in deep learning and recommendation systems by using CtR TensorSketch for polynomial feature maps (e.g., degree p > 5). The system can now capture high-order non-linear interactions in high-dimensional sparse data (e.g., user-item features, genomic sequences) without the exponential blow-up in feature space, enabling real-time inference on streaming data.

3. Improved Compact Bilinear Pooling: Replace standard TensorSketch in compact bilinear pooling (used in fine-grained visual recognition and multimodal fusion) with CtR variant. This reduces variance in gradient estimates during backpropagation, leading to faster convergence and higher accuracy in vision-language models, while cutting training time due to the same O(nnz(x) + D log D) complexity.

4. Tighter Concentration for Randomized Algorithms: Use the reduced variance to derive better (ε, δ)-approximation guarantees for randomized numerical linear algebra tasks (e.g., low-rank approximation, matrix multiplication). AI systems can now provide probabilistic error bounds with fewer samples, enabling more reliable predictions in safety-critical applications like medical diagnosis or autonomous driving.

5. Efficient Training of Over-Parameterized Networks: Apply CtR TensorSketch to approximate the neural tangent kernel (NTK) in theoretical analyses and practical training of wide neural networks. The improved variance allows for smaller sketch sizes when estimating NTK entries, reducing the memory footprint for large-scale models while preserving convergence guarantees, making it feasible to train deeper networks on limited hardware.

6. Real-Time Kernel-Based NLP and Recommender Systems: Deploy CtR TensorSketch for polynomial kernel approximations in text classification, sentiment analysis, and collaborative filtering. The input-sparsity time and reduced variance enable real-time updates on user interactions (e.g., clickstreams) with higher accuracy than existing methods, improving personalization and responsiveness in production systems.

7. Genomic and Biomedical Data Analysis: Use CtR TensorSketch to compute high-order polynomial kernels on sparse genomic feature vectors (e.g., SNP data). The reduced variance and faster runtime allow AI systems to identify complex epistatic interactions (gene-gene interactions) that were previously computationally prohibitive, enabling more accurate disease risk prediction and drug response modeling.

TensorSketch by provides efficient sketching algorithms for high-dimensional polynomial kernels p in d p. uses dense Johnson-Lindenstrauss (JL)-type projections with computational cost O(pDd), where D denotes the sketch dimension, whereas extends the sparse CountSketch algorithm, yielding a faster algorithm for high-dimensional sparse inputs with running time O (p(+ D D)). However, the variance of both estimators grows exponentially with the polynomial degree p, scaling as 3 p/D. Recent work by showed that using complex-valued distribution reduces this dependence to 2 p/D for the approach of. However, their method relies on dense JL-type projections with computational cost O(pDd) and does not extend to the algorithm of. In this work, we introduce a simple variant of TensorSketch that achieves the same variance bound as, while retaining its advantage of the input-sparsity running time. We validate our results with supporting experiments on synthetic and real-world datasets.