Language Modeling Using Tensor Trains

TL;DR

This paper introduces TTLM, a Tensor Train-based language model that represents sentences in a high-dimensional tensor space and computes probabilities through a low-dimensional contraction, unifying several RNN variants under a tensor-network framework. It formulates TT decomposition of the weight tensor, derives a recursive probability computation, and defines two practical variants, TTLM-Tiny and TTLM-Large. Empirical results on WikiText-2 and PTB show TTLM variants outperforming Vanilla-RNNs in low-scale settings and clarifying overfitting dynamics, while revealing a close relationship to Second-order RNNs, RACs, and MI-RNNs. The work demonstrates the viability of tensor-network language models for real-world data and lays groundwork for future study of normalization and long-range dependency modeling in TT-based architectures.

Abstract

We propose a novel tensor network language model based on the simplest tensor network (i.e., tensor trains), called `Tensor Train Language Model' (TTLM). TTLM represents sentences in an exponential space constructed by the tensor product of words, but computing the probabilities of sentences in a low-dimensional fashion. We demonstrate that the architectures of Second-order RNNs, Recurrent Arithmetic Circuits (RACs), and Multiplicative Integration RNNs are, essentially, special cases of TTLM. Experimental evaluations on real language modeling tasks show that the proposed variants of TTLM (i.e., TTLM-Large and TTLM-Tiny) outperform the vanilla Recurrent Neural Networks (RNNs) with low-scale of hidden units. (The code is available at https://github.com/shuishen112/tensortrainlm.)

Figures (7)

• Figure 1: A quick introduction to tensor diagram notation. There are two rules of tensor diagrams: (1) tensors are notated by solid shapes with a number of 'legs' corresponding to their indices; (2) connecting two index lines implies a contraction or summation over the connected indices. In this paper, we augment our equations with these diagrams to make them easier to understand.
• Figure 2: a) Tensor Train Language Model based on Eq. \ref{['eq: ttlm_general']}. b) TT core of TTLM-Tiny. c) TT core of TTLM-Large. The dashed line in the square represents $\mathcal{A}, \Phi(X)$, or ${\bm{\mathsfit{G}}}$. Note that the only difference between TTLM-Large and TTLM-Tiny is whether to use tensor ${\bm{\mathsfit{W}}}^{eh}$.
• Figure 3: Recursive calculation of conditional probability in TTLM. Here we provide an example that given the text $x^{(1:3)}$, ${\bm{y}}^{(4)}=\psi({\bm{\mathsfit{G}}}^{(4)}{\bm{h}}_{\textrm{TTLM}}^{(3)})$ where ${\bm{y}}^{(4)}$ is the probability distribution of word $x^{(4)}$.
• Figure 4: Test set PPL on the PTB dataset w.r.t. ranks/hidden units. RNNs here denotes Vanilla-RNNs, which has the same embedding size as TTLM-Large and TTLM-Tiny. RNNs-100, RNNs-200, and RNNs-300 are the Vanilla-RNNs with fixed embedding sizes of 100, 200, and 300, respectively.
• Figure 5: Validation set PPL of TTLM-Large and TTLM-Tiny with increasing ranks on the PTB dataset. Top: TTLM-Large, Bottom: TTLM-Tiny.

Theorems & Definitions (4)

• Claim 2.1
• proof
• Claim 2.2
• proof