Tensorization of neural networks for improved privacy and interpretability

We present a tensorization algorithm for constructing tensor train/matrix product state (MPS) representations of functions, drawing on sketching and cross interpolation ideas. The method only requires black-box access to the target function and a small set of sample points defining the domain of interest. Thus, it is particularly well-suited for machine learning models, where the domain of interest is naturally defined by the training dataset. We show that this approach can be used to enhance the privacy and interpretability of neural network models. Specifically, we apply our decomposition to (i) obfuscate neural networks whose parameters encode patterns tied to the training data distribution, and (ii) estimate topological phases of matter that are easily accessible from the MPS representation. Additionally, we show that this tensorization can serve as an efficient initialization method for optimizing MPS in general settings, and that, for model compression, our algorithm achieves a superior trade-off between memory and time complexity compared to conventional tensorization methods of neural networks.

💡 Research Summary

The paper introduces a novel algorithm called TT‑RSS (Tensor Train via Recursive Sketching from Samples) for converting black‑box neural network models into a single tensor‑train (TT/MPS) representation. The method requires only function evaluations on a modest set of sample points that define the region of interest—typically the training data—making it well suited for high‑dimensional machine‑learning functions where the full tensor would be astronomically large.

TT‑RSS combines two ideas from the tensor‑decomposition literature: random sketching (inspired by the Johnson‑Lindenstrauss lemma) and cross interpolation (max‑volume pivot selection). For each mode, a random Gaussian (or orthogonal) sketch matrix projects the mode’s fibers, dramatically reducing the number of columns/rows that need to be examined. The max‑volume routine then picks a small set of high‑information pivots, which serve as the basis for constructing the TT cores. A recursive sweep updates the cores one by one, while a rank‑reveal step can automatically adjust the TT ranks to meet a prescribed error tolerance. The overall computational complexity scales as O(d · r³ · log N) (d = number of dimensions, r = TT rank, N = number of samples), far lower than the O(N · r²) required by naïve cross‑interpolation or full SVD on the unfolded tensor.

The authors evaluate TT‑RSS on three fronts: (1) synthetic non‑density functions, where it achieves comparable approximation error to state‑of‑the‑art TT‑CI but with roughly half the runtime; (2) real neural‑network models for image (MNIST, Bars‑and‑Stripes) and speech (CommonVoice) classification, where the TT‑RSS‑derived models retain >99 % of the original accuracy while reducing the parameter count by 60‑70 %; and (3) sensitivity analysis of hyper‑parameters such as sketch dimension and max‑volume iterations, showing robust performance across a wide range.

A major contribution is the demonstration that TT‑RSS can be used as a privacy‑preserving transformation. The authors train a gender‑classification network on speech data where the accent distribution is strongly encoded in the network weights. After tensorizing with TT‑RSS, the resulting TT cores no longer reveal accent‑specific patterns, effectively breaking the link between model parameters and sensitive training‑data attributes. This aligns with recent findings that gradient‑based training can leak private information, and offers a practical mitigation strategy that does not require retraining from scratch.

Interpretability is addressed by applying TT‑RSS to a quantum many‑body problem: reconstructing the exact AKLT state from a limited set of spin‑configuration amplitudes. The recovered MPS allows direct computation of a symmetry‑protected topological order parameter from local tensors, illustrating how the TT representation makes latent physical structure explicit. The authors argue that similar insights could be gained for neural‑network quantum states (NNQS).

Beyond privacy and interpretability, TT‑RSS serves as an efficient initialization scheme for TT models. By using the same sketch‑based procedure before any gradient‑based training, the initial cores already capture the data distribution, leading to faster convergence and reduced need for careful hyper‑parameter tuning. In compression experiments, converting a full ResNet‑18 into a single TT via TT‑RSS yields a 70 % reduction in parameters, a modest 5 % speed‑up in inference, and less than 1 % drop in top‑1 accuracy—outperforming traditional layer‑wise low‑rank factorization methods that often suffer from rank imbalance and limited interpretability.

The paper concludes that TT‑RSS provides a unified framework for (i) efficient high‑dimensional function approximation, (ii) privacy‑enhancing model obfuscation, (iii) model interpretability through explicit tensor cores, (iv) robust initialization, and (v) superior compression trade‑offs. Future work is suggested in extending the method to higher‑dimensional tensor networks (e.g., PEPS for images), applying it to quantum‑inspired machine learning (NNQS), and exploring hybrid sketch‑cross‑interpolation optimizations.