Low-Rank Constraints for Fast Inference in Structured Models

When modeling complex sequential spaces, such as sentences, musical scores, or video frames, a key choice is the internal structural representations of the model. A common choice in recent years is to use neural representations (Bengio et al., 2003; Mikolov et al., 2011; Brown et al., 2020; Boulanger-Lewandowski et al., 2012; Huang et al., 2018; Weissenborn et al., 2020) to store a deterministic history. These models yield strong predictive accuracy but their deterministic, continuous forms provide little insight into the intermediate decisions of the model.^††Code is available here.

Latent structured models provide an alternative approach where complex modeling decisions are broken down into a series of probabilistic steps. Structured models provide a principled framework for reasoning about the probabilistic dependencies between decisions and for computing posterior probabilities. The structure of the decision processes and the ability to answer queries through probabilistic inference afford interpretability and controllability that are lacking in neural models (Koller and Friedman, 2009; Levine, 2018).

Despite the benefits of structured models, the computational complexity of training scales asymptotically much worse than for neural models, as inference, and therefore training, requires marginalizing over all possible latent structures. For standard general-purpose models like Hidden Markov Models (HMM) and Probabilistic Context-Free Grammars (PCFG), the runtime of inference scales quadratically and cubically in the number of states respectively, which limits the ability to reach a massive scale. Promisingly, recent work has shown that in specific situations these models can be scaled, and that the increased scale results in commensurate improvements in accuracy – without sacrificing the ability to perform exact inference (Dedieu et al., 2019; Chiu and Rush, 2020; Yang et al., 2021).

In this work, we propose an approach for improving the runtime of a large class of structured latent models by introducing a low-rank constraint. We target the family of models where inference can be formulated through a labeled directed hypergraph, which describes a broad class of dynamic-programming based inference (Klein and Manning, 2004; Huang and Chiang, 2005; Zhou et al., 2006; Javidian et al., 2020; Chiang and Riley, 2020). We show how under low-rank constraints these models allow for more efficient inference. Imposing a low-rank constraint allows for a key step of inference to be rewritten as a fast matrix-vector product. This approach is also inspired by recent advances in computationally efficient neural attention attention (Katharopoulos et al., 2020; Peng et al., 2021; Choromanski et al., 2020), a significantly different task and formulation, that rewrites matrix-vector products as fast low-rank products using approximate kernel techniques.

We evaluate this approach by learning low-rank structured models for the tasks of language modeling, polyphonic music modeling, unsupervised grammar induction, and video modeling. For these tasks we use a variety of models including HMMs, PCFGs, and Hidden Semi-Markov Models (HSMMs). As the application of low-rank constraints is nontrivial in high-dimensional structured models due to reduced expressivity, we demonstrate effective techniques for overcoming several practical challenges of low-rank parameterizations. We find that our approach achieves very similar results to unconstrained models at large state sizes, while the decomposition allows us to greatly increase the speed of inference. Results on HMMs show that we can scale to more than 16,000 states; results on PCFGs achieve a significant perplexity reduction from much larger state spaces compared to past work (Kim et al., 2019); and results on HSMMs show that our formulation enables scaling to much larger state spaces for continuous emissions (Fried et al., 2020).

We consider the problem of modeling a sequence of observations $p(x)=p(x_{1},\dots,x_{T})$. These observations can range in complexity from the words in a sentence to a series of co-occurring musical notes, or to features of video frames, and may be discrete or continuous. We assume these observations are generated by an unobserved (latent) structured representation $z$, and therefore model the joint $p(x,z)$. The structure may be sequential or hierarchical, such as latent trees, and the set of structures $\mathcal{Z}$ is combinatorial, i.e. exponential in size with respect to the input sentence. In order to train these models on observations, we must optimize the evidence $p(x)=\sum_{z}p(x,z)$ by marginalizing over $z$. Scaling this marginalization is the focus of this work.

Hypergraphs are a graphical model formalism for structured distributions that admit tractable inference through dynamic programming  (Klein and Manning, 2004; Huang and Chiang, 2005; Zhou et al., 2006; Javidian et al., 2020; Chiang and Riley, 2020).¹¹1While the formalism is similar to undirected factor graphs, it allows us to represent more complex distributions: notably dependency structures with unknown topologies, such as latent trees. A labeled, directed, acyclic hypergraph consists of a set of nodes $\cal V$, a set of hyperedges $\cal E$, and a designated root node $S\in{\cal V}$. Each node $v\in\mathcal{V}$ has a collection of labels ${\cal L}_{v}$. Each hyperedge $e\in\mathcal{E}$ has a head node $u$ and tuple of tail nodes, $v=(v_{1},\ldots,v_{|e|})$, where $|e|$ is the number of tail nodes. For simplicity, we will assume at most 2 tail nodes $v_{1},v_{2}$, and unless noted, a fixed label set $\mathcal{L}$ throughout. Each hyperedge $e$ is associated with a score matrix $\Psi_{e}\in\mathbb{R}^{\mathcal{L}\times\mathcal{L}^{|e|}}$ with a score for all head and tail labels.²²2This formalism can represent inference in both locally and globally normalized models, although we focus on local normalization in this work. We use the notation $[\Psi_{e}]_{z_{u},(z_{1},z_{2})}$ to indicate the score for head label $z_{u}$ and tail labels $z_{1}$ and $z_{2}$. Finally, we assume we have a topological ordering over the edges.

A hypergraph is used to aggregate scores bottom-up through a dynamic programming (belief propagation) algorithm. Algorithm 1 (left) shows the algorithm. It works by filling in a table vector $\alpha_{v}\in\mathbb{R}^{\mathcal{L}}$ for each node $v$ in order, and is initialized to 1 at the leaf nodes.³³3 The accumulation of scores is denoted by $\stackrel{{\scriptstyle+}}{{\leftarrow}}$. Multiple hyperedges can have the same head node, whose scores must be added together. It returns the sum over latent structures, $p(x)$. Counting loops, the worst-case runtime complexity is $O(|{\cal E}|\times L^{|e^{*}|+1})$ where $L=|{\cal L}|$ is the size of the label set and $|e^{*}|$ the max hyperedge tail size. Algorithm 1 (right) shows the same algorithm in matrix form by introducing joined tail vectors $\beta_{v}\in\mathbb{R}^{\mathcal{L}^{|e|}}$ for each group of nodes $v$. Letting $z_{v}=(z_{1},z_{2})$, the joined tail vector contains entries $[\beta]_{z_{v}}=[\alpha_{v_{1}}]_{z_{1}}[\alpha_{v_{2}}]_{z_{2}}$.

To make this formalism more concrete, we show how hypergraphs can be used for inference in several structured generative models: hidden Markov models, probabilistic context-free grammars, and hidden semi-Markov models. Inference in these examples are instances of the hypergraph algorithm.

Example: Hidden Markov Models (HMM) HMMs are discrete latent sequence models defined by the following generative process: first, a sequence of discrete latent states $z=(z_{1},\ldots,z_{T})$ with state size $L$ are sampled as a Markov chain. Then each state $z_{t}$ independently emits an observation $x_{t}$, i.e.

where $p(z_{t}\mid z_{t-1})$ is the transition distribution, $p(x_{t}\mid z_{t})$ the emission distribution, and $p(z_{1}\mid z_{0})$ is the initial distribution with distinguished start symbol $z_{0}$.

Given a sequence of observations $x=(x_{1},\ldots,x_{n})$ we can compute $p(x)=\sum_{z}p(x,z)$ using a labeled directed hypergraph, with single-tailed edges, nodes corresponding to state positions, labels corresponding to states, and emissions probabilities incorporated into the scoring matrices $\Psi$. There are $T$ scoring matrices, $\Psi_{t}\in\mathbb{R}^{\mathcal{L}\times\mathcal{L}}$, with entries $[\Psi_{t}]_{z_{t},z_{t+1}}=p(z_{t+1},x_{t}\mid z_{t})$ corresponding to transitions.⁴⁴4 The left-most scoring matrix for the HMM has entries $[\Psi_{1}]_{z_{1},z_{2}}=p(z_{2},x_{1}\mid z_{1})p(z_{1}\mid z_{0})$. Algorithm 2 (left) shows the approach. This requires time $O(TL^{2})$ and is identical to the backward algorithm for HMMs.⁵⁵5 In the case of HMMs, the table vectors $\alpha_{t}$ correspond to the backward algorithm’s $\beta$ values.

Example: Context-Free Grammars (CFG) CFGs are a structured model defined by the 5-tuple $\mathcal{G}=(S,\mathcal{N},\mathcal{P},\mathcal{X},\mathcal{R})$, where $S$ is the distinguished start symbol, $\mathcal{N}$ is a set of nonterminals, $\mathcal{P}$ is a set of preterminals, $\mathcal{X}$ is the token types in the vocabulary, and $\mathcal{R}$ is a set of grammar rules. Production rules for start, nonterminals, and preterminals take the following forms:⁶⁶6 We restrict our attention to grammars in Chomsky normal form.

A probabilistic context-free grammar (PCFG) additionally has a probability measure on the set of rules. To compute $p(x_{1},\ldots,x_{T})$ with a hypergraph, we create one node for each contiguous subspan $[i,k)$ in the sentence. Nodes with $i+1<k$ have a nonterminal label set ${\cal L}=\mathcal{N}$. Nodes with $i+1=k$ have a preterminal label set ${\cal L}_{i,i+1}=\mathcal{P}$. The main scoring matrix is $\Psi\in\mathbb{R}^{\mathcal{L}\times\mathcal{L}^{2}}$, with entries $[\Psi]_{z_{u},(z_{1},z_{2})}=p(z_{1},z_{2}\mid z_{u})$.⁷⁷7We have a separate matrix for terminal production on $x$ which we elide for simplicity. Algorithm 2 (right) shows how for every hyperedge we join the scores from the two tail nodes in $\alpha_{i,j}$ and $\alpha_{j,k}$ into joined tail vector $\beta_{i,j,k}\in\mathbb{R}^{\mathcal{L}^{2}}$. As there are $O(T^{3})$ hyperedges and the largest ${\cal L}$ is of size $|\mathcal{N}|$, the runtime of the algorithm is $O(T^{3}{|\mathcal{N}|}^{3})$. This approach is identical to the CKY algorithm.

Example: Hidden Semi-Markov Models (HSMM) HSMMs are extensions of HMMs that allow for generating a variable-length sequence of observations per state. It defines the following generate process: first, we sample a sequence of discrete latent states $z=(z_{1},\cdots,z_{K})$ with a first-order Markov model. We then use them to generate the length of observations per state. For our experiments we generate independent continuous emissions $x_{t}$ with a Gaussian distribution for $p(x_{i}\mid z_{k})$. Full details of the inference procedure are given in Appendix E.

For these structured distributions, hypergraphs provide a general method for inference (and therefore training parameterized versions). However, the underlying algorithms scale poorly with the size of the label sets (quadratic for HMM and HSMM, cubic for CFG). This complexity makes it challenging to scale these models and train versions with very large numbers of states.

In this section, we consider an approach for improving the scalability of these models by reducing the dependence of the computational complexity of inference on the label set size. The main idea is to speed up the matrix-vector product step in inference by using a low-rank decomposition of the scoring matrix $\Psi$. In the next section we show that this constraint can be easily incorporated into parameterized versions of these models.

We evaluate the application of low-rank constraints with four experiments: sequential language modeling with HMMs, polyphonic music modeling with a large observation space, hierarchical language modelings with PCFGs, and video modeling with HSMMs.

Data Our first set of experiments evaluate sequential models on Penn Treebank dataset (Ptb) (Marcus et al., 1993) for the task of word-level language modeling. We use the preprocessing from Mikolov et al. (2011). The second set of experiments is on polyphonic music modeling (Boulanger-Lewandowski et al., 2012). We evaluate on four music datasets: Nottingham (Nott), Piano, MuseData (Muse), and JSB chorales (JSB). Each timestep consists of an 88-dimensional binary vector indicating whether a particular note is played. Since multiple notes may be played at the same time, the effective vocabulary size is extremely large. The third set of experiments use PCFGs for language modeling, we also use Ptb, but with the splits and preprocessing used in unsupervised constituency parsing (Shen et al., 2018, 2019; Kim et al., 2019). The last set of experiments use HSMMs for video modeling, where we use CrossTask (Zhukov et al., 2019) with 10% of the training data for validation. We follow the preprocessing steps in Fried et al. (2020) and apply PCA to project features to vectors of size 200. For the full details on datasets, please see Appendix D.

Models and Hyperparameters For language modeling with HMMs, we experiment with a range of state sizes, $|\mathcal{L}|=L\in\left\{2^{10},2^{11},2^{12},2^{13},2^{14}\right\}$, and rank $N\in\left\{L/2,L/4,L/8\right\}$. For polyphonic music modeling with HMMs, we experiment with states sizes $L\in\left\{2^{7},2^{8},2^{9},2^{10},2^{11}\right\}$. For language modeling with PCFGs, we use a set of nonterminals of size $|\mathcal{N}|\in\left\{30,60,100\right\}$ and preterminals of twice the number of nonterminals $|\mathcal{P}|=2|\mathcal{N}|$. Our smallest setting ($|\mathcal{N}|=30$, $|\mathcal{P}|=60$) is the one used in Kim et al. (2019). For video modeling with HSMMs, we use the same model setting as Fried et al. (2020), but we don’t constrain states to the predefined states per task, and we experiment with state sizes $L\in\left\{2^{6},2^{7},2^{8},2^{9},2^{10}\right\}$ and rank $N\in\left\{2^{4},2^{5},2^{6},2^{7}\right\}$.

We utilize the feature map $\phi(x)=\exp(Wx)$ for the LHMM and LHSMM, and $\phi(x)=\exp(Wx-\|x\|_{2}^{2}/2)$ for the LPCFG. We initialize the parameters of feature maps using orthogonal feature projections (Choromanski et al., 2020), and update it alongside the model parameters. For the full hyperparameter and optimization details, see Appendix G.

Baselines and Evaluation The language modeling experiments are evaluated using perplexity. Baselines are neurally parameterized HMM with a standard softmax transition. We also compare to VL-HMM, which makes a strong structural sparsity assumption on the emission distribution (Chiu and Rush, 2020). We include for reference a state-of-the-art language model, the AWD-LSTM (Merity et al., 2017). For polyphonic music modeling, we compare our LHMM against RNN-NADE (Boulanger-Lewandowski et al., 2012) which models the full joint distribution of notes as well as temporal dependencies; as well as autoregressive neural models such as the R-Transformer (Wang et al., 2019) (as reported by Song et al. (2019)) and an LSTM  (as reported by Ziegler and Rush (2019)); models with latent continuous dynamics such as the LV-RNN (Gu et al., 2015) and SRNN (Fraccaro et al., 2016); and finally comparable models with latent discrete dynamics, the TSBN (Gan et al., 2015) and the baseline HMM. We evaluate perplexities of our low-rank PCFG (LPCFG) against a softmax PCFG (PCFG) (Kim et al., 2019). For video modeling, we evaluate the negative log likelihoods on the test set and compare low-rank HSMMs to softmax HSMMs.

Hidden Markov Models for Language Modeling

Our main experimental result is that the low-rank models achieve similar accuracy, as measured by perplexity, as our baselines. Fig. 1 shows that perplexity improves as we increase the scale of the HMM, and that the performance of our LHMM also improves at the same rate. At small sizes, the low-rank constraints slightly hinder accuracy; however once the size is large enough, i.e. larger than $2^{12}$, LHMMs with 8:1 state-to-rank ratios perform comparably. ¹⁰¹⁰10 See Appendix H for an analysis of the ranks of HMMs/LHMMs.

Fig. 1 also contains speed comparisons between HMMs and LHMMs. A state-to-rank ratio of 8:1 matches the accuracy of softmax HMMs at larger state sizes and also gives an empirical speedup of more than 3x at $L=2^{14}$. As expected, we only see a speedup when the state-to-rank ratio exceeds 2:1, as we replaced the $O(L^{2})$ operation with two $O(LN)$ ones. This implies that the low-rank constraint is most effective with scale, where we observe large computational gains at no cost in accuracy.

HMMs are outperformed by neural models, and also by VL-HMMs (Chiu and Rush, 2020) which offer similar modeling advantages to HMMs, as shown in Tbl. 1 (left). This indicates that some aspects of performance are not strictly tied to scale. We posit this is due to the problem-specific block-sparse emission constraint in VL-HMMs. While very effective for language modeling, the VL-HMM relies on a hard clustering of states for constraining emissions. This is difficult to apply to problems with richer emission models (as in music and video modeling).

Hidden Markov Models for Music Modeling We next apply LHMMs on polyphonic music modeling. This has a max effective vocabulary size of $2^{88}$, as multiple notes may occur simultaneously. Unlike for language modeling, we use a factored Bernoulli emission model, modeling the presence of each note independently. Fig. 2 (right) shows that HMMs are competitive with many of the models on these datasets, including LSTMs. We find that LHMMs achieve performance slightly worse than but comparable to the unconstrained HMMs overall. Fig. 2 (left) shows that the distinction drops with more states. Both HMMs achieve low negative likelihoods (NLL) on the datasets with shorter sequences, Nottingham and JSB, but relatively poorer NLLs on the datasets with longer sequences (Muse and Piano).

Context-Free Grammars For syntactic language modeling on Ptb, our low-rank PCFG (LPCFG) achieves similar performance to PCFGs, as shown in Table 2 (left), with an improvement in computational complexity. The complexity of inference in PCFGs models is cubic in the number of nonterminals, so even models with $|\mathcal{N}|=30$ nonterminals are relatively costly. Our approach achieves comparable results with $N=8$ features. As we scale up the number of nonterminals to $|\mathcal{N}|=100$, LPCFG stays competitive with a lower computational complexity (since $N<|\mathcal{N}|$). These experiments also demonstrate the importance of scale in syntactic language models with more than 50 point gain in perplexity over a strong starting model.

CFG Speed Once the model is large enough, i.e. $|\mathcal{N}|\geq 60$ nonterminals and $|\mathcal{P}|\geq 120$ preterminals, the LPCFG is faster than PCFG, as shown in Tbl. 2 (left). Note that the LPCFG is faster than the CFG even when the number of features $N>\frac{|\mathcal{N}|}{2}$, in contrast to the HMM case where a speedup can only be obtained when $N<L/2$. This is due to the scoring matrix being rectangular: Recall the low-rank matrix product $\Psi\beta=U(V^{\top}\beta)$, where, when specialized to PCFGs, the left-hand side takes time $O(L^{3})$ and the right-hand side takes $O(L^{2}N+LN)$. For PCFGs, the term $V^{\top}\beta$ dominates the runtime. This contrasts with HMMs, where both $V^{\top}\beta$ and the subsequent multiplication by $U$ take the same amount of time, $O(LN)$.

Hidden Semi-Markov Models for Video Modeling Table 2 (right) shows the results of video modeling using HSMMs. In addition to using a different hypergraph for inference, these experiments use a continuous Gaussian emission model. By removing the state constraints from tasks, our HSMM baselines get better video-level NLLs than that from Fried et al. (2020) at the cost of more memory consumption. Due to GPU memory constraints, we can only train HSMMs up to $2^{8}$ states. However, the low-rank parameterization allows models to scale to $2^{10}$ states, yielding an improvement in NLL. Absolute results could likely be improved with more states and by an improved emission parameterization for all models.

Improving Rank Assumptions One potential limitation of all low-rank models is that they cannot learn high-rank structures with low $N$. We began to see this issue at a ratio of 16:1 states to features for large HMMs. To explore the effects of this limitation, we perform an additional experiment that combines low-rank features with a sparse component. Specifically we add an efficient high-rank sparse banded transition matrix. The full details are in Appendix I. Tbl. 1 (right) shows that combination with the band structure allows for larger ratios than just the low-rank structure alone, while only adding another operation that costs $O(LN)$.

Similar to our work, other approaches target matrix or tensor operations in inference, and impose structural model constraints to improve computational complexity. Many of the works on HMMs in particular take advantage of the transition structure. The Dense-mostly-constant (DMC) HMM assigns a subset of learnable parameters per row of the transition matrix and sets the rest to a constant, leading to a sub-quadratic runtime (Siddiqi and Moore, 2005). Other structures have also been explored, such as aligning the states of an HMM to underlying phenomena that allows inference to be sped up (Felzenszwalb et al., 2004; Roweis, 2000). Additionally, other methods take advantage of emission structure in HMMs in order to scale, such as the Cloned HMM (Dedieu et al., 2019) and VL-HMM (Chiu and Rush, 2020). Compared to these approaches, our method is more flexible and generic, since it can be applied in a non-application-specific manner, and even extended with high-rank components (such as banded structure).

Low-rank structure has been explored in both HMMs (Siddiqi et al., 2009), a generalization of PCFGs called weighted tree automata (Rabusseau et al., 2015), and conditional random fields (Thai et al., 2018). The reduced-rank HMM (Siddiqi et al., 2009) has at most 50 states, and relies on spectral methods for training. The low-rank weighted tree automata (Rabusseau et al., 2015) also trains latent tree models via spectral methods. We extend the low-rank assumption to neural parameterizations, which have been shown to be effective for generalization (Kim et al., 2019; Chiu and Rush, 2020), and directly optimize the evidence via gradient descent. Finally, Thai et al. (2018) do not take advantage of the low-rank parameterization of their CRF potentials for faster inference via low-rank matrix products, a missed opportunity. Instead, the low-rank parameterization is used only as a regularizer, with the full potentials instantiated during inference.

Concurrent work in unsupervised parsing uses a tensor decomposition to scale PCFGs to large state spaces (Yang et al., 2021). Our low-rank decomposition of the flattened head-tail scoring matrix is more general, resulting in worse scaling for the PCFG setting but with wider applicability, as shown by experiments with HMMs and HSMMs.

This work improves the scaling of structured models by establishing the effectiveness of low-rank constraints for hypergraph models. We show that viewing a key step of inference in structured models as a matrix-vector product, in combination with a low-rank constraint on relevant parameters, allows for an immediate speedup. Low-rank inference allows us to obtain a reduction in the asymptotic complexity of marginalization at the cost of a constrained model. Our approach applies to a wide class of models, including HMMs, HSMMs, and PCFGs. Through our experiments on language, video, and polyphonic music modeling, we demonstrate an effective approach for overcoming the practical difficulty of applying low-rank constraints in high dimensional, structured spaces by targeting and constraining model components that bottleneck computation. Future work includes exploration of other structural constraints for speeding up matrix-vector products (Dao et al., 2020) performed in inference, as well as application to models where exact inference is intractable.

We thank Nikita Kitaev for the discussion that sparked this project. We thank Sam Wiseman, Jack Morris, and the anonymous reviewers for valuable feedback. Yuntian Deng is sponsored by NSF 1704834, Alexander Rush by NSF CAREER 1845664, and Justin Chiu by an Amazon research award.