Improving TensorSketch Using Complex Random Variables

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

Any-Dimensional Learning by Sampling

Many machine learning models are defined for inputs of different sizes, such as point clouds containing different numbers of points, sequences of tokens of different lengths, and graphs on different numbers of nodes. Such models are trained on finitely-many examples of necessarily limited sizes. How well do these models generalize from inputs of small size to larger inputs of size not seen during training? Furthermore, evaluating such models on large inputs is often expensive. How can we sketch large inputs to obtain smaller ones on which the model takes similar values? At the heart of both questions is the need to compare inputs of different sizes and to approximate large inputs by small ones. We present a unified approach to address these questions by using random sampling maps to compare inputs of different sizes. The sampling maps we consider are generalizations of sampling with replacement, random binning, and species sampling. We characterize the application domains in which each type of sampling is appropriate in terms of the symmetries and relations between problem instances of different sizes in the domain. Our framework yields explicit generalization and sketching rates for function classes continuous with respect to a chosen notion of sampling, encompassing large families of functions defined on sequences, graphs, and tensors of different sizes. Specific examples include moment polynomials on measures, homomorphism densities and numbers of graphs, permutation-invariant transformers, and graph neural networks.

Stochastic trace estimation with tensor train random vectors

Stochastic trace estimation is a standard tool for approximating the trace of a large-scale matrix available only through matrix-vector products. However, in tensor-structured settings, unstructured Gaussian or Rademacher test vectors may be prohibitively expensive to store and compute with, while cheaper rank-one tensor-product vectors can require sample complexities that grow exponentially with the tensor order. This work studies Gaussian random tensor train vectors as a structured alternative for stochastic trace estimation. We show that, with a suitable choice of the tensor train rank, random tensor train vectors recover dimension-independent guarantees for the Girard--Hutchinson estimator. In particular, a median-of-means variant with tensor train rank $r \geq d-1$ achieves the same dependence on the accuracy $\varepsilon$ and failure probability $δ$ as the classical estimator based on unstructured Gaussian vectors. We further prove an oblivious subspace injection result for sketches formed from independent Gaussian random tensor train vectors: tensor train rank $r\geq d-1$ and $\mathcal{O}(\varepsilon^{-2}(k+\log(1/δ)))$ samples suffice for a $k$-dimensional target subspace. Finally, we investigate the use of such sketches within the Nyström++ framework. We show that the resulting estimator can achieve the desired $\mathcal{O}(\varepsilon^{-1})$ sample complexity under an additional spectral-tail condition. These results provide clarififcation on both the potential and the limitations of random tensor train vectors in stochastic trace estimation.