Compressed Nonparametric Language Modelling

1 Introduction

Statistical Language Models (LM) are the key components of many tasks in Natural Language Processing, such as Statistical Machine Translation, and Speech Recognition. Conventional LMs are n-gram models which apply the Markov assumption to approximate the probability of a sequence wN 1 as,

Several smoothing techniques have been proposed to address the statistical sparsity issue in computation of each conditional probability term. The widely used smoothing techniques for LM are Kneser-Ney (KN) [Kneser and Ney, 1995], and its extension Modified Kneser-Ney (MKN) [Chen and Goodman, 1999]. The intuition behind KN, MKN and their extensions [Shareghi et al., 2016a] is to adjust the original distribution to assign non-zero probability to unseen or rare events. This is achieved by re-allocating the probability mass in an interpolative procedure via absolute discounting.

It turns out that the Bayesian generalisation of KN family of smoothing is the Hierarchical Pitman-Yor Process (HPYP) LM [Teh, 2006a], which was originally developed for finiteorder LM [Teh, 2006b], and was extended as the Sequence Memoizer (SM) [Wood et al., 2011] to model infinite-order LMs. While capturing the long range dependency via HPYP improves the estimation of conditional probabilities, these types of models remain impractical due to several computational and learning challenges, namely large model size (data structure representing the model, and the number of parameters), long training and test time, and poor sampler mixing.

In this paper we address aforementioned issues; inspired by the recent advances in using compressed data structures in LM [Shareghi et al., 2015; 2016b] our model is built on top of a compressed suffix tree (CST) [Ohlebusch et al., 2010]. In the training step, only the CST representation of text is constructed, allowing for a very fast training, while proposing an efficient approximate inference algorithm for the test time. Mixing issue is avoided via careful sampler initialisation and design.

The empirical results show that our proposed approximation of HPYP is richer than KN and MKN, and is much more efficient in learning and inference phase compared to full HPYP. Compared with 10-gram KN and MKN models, our ∞-gram model consistently improves the perplexity by up to 15%. Our compressed framework allows us to train on large collection of text, i.e. 100× larger than the largest dataset used in HPYP LMs [Wood et al., 2011] while having several orders of magnitudes smaller memory footprint and supporting fast and efficient inference.

2 Interpolative Language Models

Conventional interpolative smoothing techniques in LM follow a general form,

where, u = wi-1 i-n+1 is called the context, π(u) is u with its least recent symbol dropped, and c and d are the count and absolute discount, while γ is the mass allocated to the lower level of the interpolation. Interpolative smoothing assumes that P(w|u), the conditional distribution of a word w in the context u, is similar to and hence smoothed by that of suffix of the context P(w|π(u)). This recursive smoothing stops at the unigram level where the conditioning context is empty, u = ε (see the left panel of Table 1). In what follows, we provide a brief overview of hierarchical Pitmon-Yor Process LM and its relationship with KN.

2.1 Hierarchical Pitman-Yor Process (HPYP) LM

We start by describing the Pitman-Yor process (PYP; [Pitman and Yor, 1997]), used as a prior over LM parameters. PYP(d, θ,H) is a distribution over the space of probability distributions with three parameters: a base distribution H which is the expected value of a draw from a PYP, the concentration parameter -d < θ which controls the variation of draws from PYP around H, and the discount parameter 0 ≤ d < 1 which controls the heavy-tailness of the sampled distributions. PYP is an appropriate prior for LMs to capture power-law behaviour prevalent in the natural language [Goldwater et al., 2011].

To illustate the use of the PYP prior, we conside as the likelihood a simple unigram LM G, from which the words of a text are generated. The Chinese restaurant process (CRP) is a metaphor that allows generating words from a PYP without directly dealing with the LM G itself by integrating it out. Consider a restaurant where customers are seated on different tables, and each table is served one dish. To make the analogy, the restaurant corresponds to the LM, the customers are the text token needed to be generated, and the dishes are the words. Let tw denote the number of tables serving the same dish w in the restaurant, nw denote the total number of customers seated on these tables, and t. = ∑w tw and n. = ∑w nw to be the total number of tables and customers, respectively. Generating the next word from the LM is done by sending a customer to the restaurant which either (i) sits on an existing table serving the dish w with probability proportional to nw - dtw, or (2) sits on a new table with probability proportional to θ + dt. and orders a dish from the base distribution H. The probability of the next word w is thus

where η = {d, θ, {nw, tw}w∈σ}, and σ is the vocabulary. Note that {nw, tw}w∈σ form the sufficient statistics for generating the next word.

In a HPYP LM, the distribution Gu of words following a context u has a PYP(du, θu, Gπ(u) ) prior,

where the base distribution Gπ(u) itself has a PYP prior. This induces a hierarchy among these context-conditioned distributions (see Figure 1) tying the base distribution of each node of the hierarchy to the distribution of its parent. Note that u refers to both a context (a sequence of words) and its corresponding node in the HPYP tree. The distribution at the root of the hierarchy Gε corresponds to the empty context ε:

where the base distribution U is the uniform distribution over the vocabulary σ.

Having a HPYP LM, the next word is generated by integrating out all of the distributions corresponding to the nodes of the tree. This is achieved by hierarchical Chinese restaurant process (HCRP), whereby a customer is sent to a restaurant of a context from which a word needs to be generated. In case the customer is seated on a new table, a new customer is sent to the parent restaurant for ordering the dish. The number of customers sent to the parent is orchestrated via the concentration and discount parameters. The following constraints hold for {nuw , tuw }w∈σu across the tree nodes:

(a)

Figure 2: (top) Suffix Tree, (bottom) Suffix Array and BurrowsWheeler Transformation for "#abc#bbc#$". In (top), labels correspond to concatenation of edge labels from the root to each node and digits stored on leaves and in SA, (bottom), correspond to the starting index of the suffixes in text.

where σu is the subset of σ observed after the context u. Tying the distributions as a hierarchy simulates the smoothing process by adjusting the base distributions of each level recursively. Given {ηu}u∈HPYP , u ∈ HPYP denoting all nodes in HPYP, the predictive probability of a word w in a context u is computed recursively as PHPYP(w|u) in Table 1.

Note the close similarity in formulation of KN smoothing and the HPYP as illustrated in Table 1; indeed we recover KN smoothing from HPYP by assuming {θu = 0}u∈HPYP and {tuw = 1}u∈HPYP w∈σu . In what follows, we exploit this property to compactly represent the HPYP LM.

3 Compressed HPYP LM

The child-parent relationship between distributions in HPYP forms a context tree which can be represented as a suffix trie, or more compactly as a suffix tree [Wood et al., 2011]. In practice suffix trees require at best 20|T | bytes of space, where T denotes the text, making them impractical for anything but small data.

Hence, even storing the structure of the HPYP model without storing parameters {ηu}u∈HPYP is impractical for large datasets. This was possibly the case given that the largest dataset that HPYP LMs were built on small corpora including only a few million words [Teh, 2006b; Gasthaus et al., 2010; Wood et al., 2011]. However, with the availability of datasets of orders of magnitudes larger size [Buck et al., 2014; Parker et al., 2011], it is crucial to improve the scalability of HPYP LMs. In the following, we briefly introduce a set of compressed data structures and operations, and then illustrate how a HPYP LM can be made scalable using these tools.

3.1 Compressed Suffix Tree

A Suffix Tree (ST) [Weiner, 1973] of a string T with alphabet σ is a tree of |T | + 1 leaves, where a path from the root to a leaf corresponds to a suffix of T . Each leaf holds a number indicating the staring position of the suffix in T while leaves are lexicographically ordered, see Figure 2(top). The search for any sequence u in T corresponds to finding the node v in ST such that u is a prefix of the concatenation of the path labels from the root to v. While suffix trees offer O(|u|) search complexity, in practice they require 20|T | bytes.

A Suffix Array (SA) [Manber and Myers, 1993] of T is an array of sorted suffixes of T , where SA[i] holds a same number as i-th leaf in ST, see Figure 2(bottom). Search in SA translates to binary search to find the corresponding range that spans over all substrings that have u as their prefix, and is O(|u| log |T |). Constructing SA takes 4-8|T | bytes in practice, which compared to |T | log |σ| bits required to store T makes both ST and SA impractical to use for large data.

A Compressed Suffix Array (CSA) exploits text compressibility and provides the same functionality as SA but in space equal to bzip2 compressed T in practice. We use the FM-Index [Ferragina et al., 2008] that utilizes the text compressibility by using the Burrows-Wheeler transformation (BWT) [Burrows and Wheeler, 1994] of the text, which is defined as BWT[i] = [SA[i] - 1 mod |T |] as illustrated in Figure 2(bottom). Searching for a sequence in BWT is done in reverse order (called backward-search) and requires O(|u| log |σ|). Similarly, a Compressed Suffix Tree (CST) simulates ST and is built on CSA by storing extra bits to store the shape of the tree and path labels [Ohlebusch et al., 2010]. For details see [Shareghi et al., 2015].

Table 2 illustrates the key operations on CSA and CST along with their complexities. The backward-search looksup the CSA span covering the given sequence and returns the [lb, rb] of the node v matching it, and the size operation counts the number of leaves of the CST subtree rooted at v using the length of the returned range. The degree operation returns the number of types completing a sequence to its right. The most expensive operation is the interval-symbols, which for a context π(u) returns its corresponding set of children and each child's corresponding node's range using the BWT [Schnattinger et al., 2010]. Figure 2(bottom) shows the interval symbol procedure for finding the children of "π(u) = bc": First a search is done to find the corresponding range for "bc" in SA, highlighted in gray, then the corresponding cells on BWT are looked up to find {b, a} as the two possible completions of "bc" to its left, and then their corresponding nodes are located efficiently, as illustrated by dashed arrows.

3.2 Compressed HPYP LM

We make use of a CST to represent HPYP LM compactly, resulting in a model which is less than the size of the text itself. This is inspired by the use of CSTs to represent KN family of smoothing, and the fact that the structure of the hierarchy representing both models are exactly the same. The basic idea for compressed KN-smoothed LMs is to extract the required counts directly from a CST

representation of the text on-the-fly [Shareghi et al., 2015; 2016b]; Table 2 covers the required quantities together with the CST operations and their time complexities.

The avid readers might notice that the key difference between KN and HPYP is in the {nuw , tuw }u∈HPYP w∈σu . While it is possible to store vectors {nuw , tuw }w∈σu for each node u of HPYP, it adds a significant load to the memory usage. This is one of the key computational issues of the inference approach taken in sequence memoizer [Wood et al., 2011], which samples these vectors in the training time and stores them before using them in testing. We address the issue by moving the sampling to the test time, hence skipping the need to store these samples.

To complete our model description, we need to cover the discount and concentration parameters. In our model, the discount parameters are set to Kneser-Ney discounts and tied based on the context size |u|, while each distribution uses its own separate concentration parameter. Our decision for fixing the discount parameters was to avoid the cost of sampling them during the inference. Also, the range for the discount parameters are very fine-grained making the gain in sampling them very negligible 1 . We show in the inference section how fixing the discounts allows us to develop an efficient sampler.

4 Fast Approximate Inference for HPYP LM

The inference in a HPYP LM translates into computing the predictive probability of a word w, in the context u. This corresponds to integrating out all the latent prior distributions and is defined as the following intractable integral,

and is approximated using samples for η. Here η = {ηu}u∈HPYP , and P(w|u, η) is defined as in Table 1 (right panel). While there are different approaches to generate samples, they are designed to run in the training phase and require explicit storage of sampled quantities. This amounts to a significant memory load which we avoid by skipping it. Consequently, it is required to generate the samples in the test phase, but the existing sampling algorithms [Gasthaus and Teh, 2010] generate samples across all nodes in HPYP which is too slow to fit our purpose. Instead, we present a novel sampler, designed for the query phase of LM, which results in a much lower memory usage, is fast, and avoids mixing issues inherited in the full samplers.

Let us assume the path from the root to the node matching the context u of a given query "uwi" and denote the set of distributions along the path by γ+, and all the other nodes of the HPYP by γ-. For example, for the query P(w = c|u = ab), γ+ = {Gab , Gb , Gε } and γ- = {Gu|u ∈ HPYP∧Gu /∈ γ+}. Then, the three computational solutions involved in our proposed approximate inference scheme are:

Branch Sampling While a full HPYP sampler, i.e. in SM [Wood et al., 2011], samples on γ+ ∪γ-, in here only the γ+ distributions and only the type matching the query word wi are selected for sampling, fixing γ- at their initialization which is KN. This means at any given state of sampling, we have {tuw = 1}u∈γw∈σu and allows for a fast inference in the test time while reducing the size of the sampling space exponentially hence reducing the risk of poor mixing.

1 Our analysis shows no improvements of perplexities when discounts were sampled, while it made the inference step slower.

Forgetting Samples Samples on γ+ are generated during the test phase, used for approximating the predictive probability and then forgotten immediately. In this process, for any given query the state of HPYP will be set to KN at the beginning of the sampling process. This keeps the memory usage of the training and inference phase close, and roughly matching the size of the compressed text.

Range Shrinking The key quantities in the sampling phase are tuw . In practice, the range 0 < tuw ≤ nuw can potentially be very large, making the sampling very slow. Instead, we follow a non-uniform sampling by shrinking the range to 1 ≤ tuw ≤ min{M, nuw } (Here M = 10). The motivation here is based on the key difference between KN and MKN which is mainly in the discount range. In MKN, which typically outperforms KN, discounts are larger than 1 and our empirical analysis on several datasizes and languages illustrate the range in practice is [0, 3]. 2 This effect to some degree is replicated in HPYP when the discounts 0 ≤ du < 1 are multiplied by tuw . Shrinking the sampling range keeps the HPYP distributions at each level of HPYP close to their corresponding KN (and MKN) counterparts, while the concentration parameter allows the distributions to have more flexibility in capturing the desired distribution.

The first two components allow fast inference while keeping the memory usage of our approach to be several orders of magnitudes smaller than the SM, hence making our approach computationally practical for large data regime. The third component seeks to take advantage of the best component of MKN smoothing, while avoiding the mixing issue that occurs for the infinite HPYP case. In the following two subsections, we provide the statistical underpinning of the sampling, and show how it can be done under CST mechanics.

4.1 Sampling

We sample using the joint distribution P({ηu}u∈HPYP ),

where (a|b)c is the Pochhammer 3 symbol, and Sd(n, t) is the generalized Stirling number of kind (-1, -d, 0) [Hsu and Shiue, 1998].

The joint distribution in eqn. 5 allows efficient sampling for tuw and nuw in the hierarchy, starting from the data level and going up in the hierarchy. The only expensive computation is for the Stirling numbers which are cached as KN discounts are used. We use the exact recursive formulation of Stirling numbers [Buntine and Hutter, 2012] and switch to asymptotic approximation 4 when t or n are large, i.e. ≥ 8000.

2 In theory, the MKN discounts can be as large as c(uwi).

3 (a|b)c = a(a + 1 × b)...(a + (c - 1) × b)

4 Using the Stirling's approximation for factorials.

Algorithm 1 Gibbs Sampler for η ∈ γ+

1: function SAMPLER(w, k, n, γ+, ⃗S,M) 2: u∗ ← γ+ k , π(u∗ ) ← γ+ k-1 3: tu∗ w ← 0, nu∗ w ←n 4: ⃗Q ← null 5: if nu∗ w ̸= 1 then 6: while tu∗ w ≤min{M, nu∗ w } do 7: if u∗ ̸= ε then 8: F ∝ P(tu∗ w |...) ▷ eqn.7 9: else 10: F ∝ P (tεw |...) ▷ eqn.8 11: ⃗Q ← (F, tu∗ w ) 12: tu∗ w ← tu∗ w +1 13: tu∗ w ← sample from( ⃗Q) 14: else 15: tu∗ w ← 1 16: θu∗ ← sample θu 17: ⃗S ← (u∗ , nu∗ w , tu∗ w , θu∗ ) 18: if u∗ ̸= ε then 19: nπ(u∗ ) w ← update(nπ(u∗ ) w , tu∗ w ) ▷ eqn.6 20: SAMPLER(w, k - 1, nπ(u∗ ) w , γ+, ⃗S,M) 21: return(⃗S)

For each Gu ∈ γ+, except the leaf level 5 , the nuw 's will be sampled jointly as tψ(u) w 's are sampled, where ψ(u) ∈ children(u). Starting from the leaf level of the hierarchy, the nuw 's are read from the data, hence fixed and tuw 's are sampled while satisfying the constraints in eqn.2, and eqn.3. Given a sampled tu∗ w at the leaf level u∗ , the nπ(u∗ ) w is updated as,

The conditional probability of the sampled tu∗ w from eqn. 5 for the non-root levels while fixing all the independent variables, P(tu∗ w |...), is proportional to,

Given sampled tu∗ w , nu∗ w for a context u, the concentration parameter θu∗ is sampled via auxiliary variables [Teh et al., 2012] using a Gamma(a,b) prior. Algorithm 1 illustrates the sampling procedure for collecting a single set of samples along γ+. The algorithm starts from the leaf level and moves up on the γ+ branch, sampling tuw and nπ(u) w jointly. The index k denotes the level on the extracted branch, and matches the k in Table1. Given a query, this process is repeated multiple times along the γ+ branch.

5 The leaf level is where the data is observed (first row of Table 1).

4.2 Sampling under CSTmechanics

Sampling under CST involves two passes in the opposite directions, see Figure 3. Given a query, one pass starts from the last word of the query and grows the pattern to its left one word at a time. This pass collects all the required nodes in γ+, and identifies the required fragmentations, while allowing to reuse spans during the backward-search, instead of operating a fresh search over the full CST span. Once all the required nodes are extracted, a second pass in the opposite direction on the γ+ branch, samples t, n, and θ. The core of sampler relies on eqn. 7 and eqn. 8, which we describe in below. Given a sampled tu∗ w for a context u∗ and word w, tu∗ . is defined as (N1+(u∗ ·) - 1) + (tu∗ w ) which assumes a table per word type for all the words occurring after u except for the word w for which tu∗ w is sampled. This translates into a degree(u) call to compute N1+(u∗ ·) and is done in constant time. Computing the other main quantities is more expensive and involves interval-symbols operation. The children nodes of u∗ , and π(u∗ )w are extracted via interval-symbols. Then, their table and count statistics are computed,

where for each ψ, except for u∗ , a degree operation is called, and tu∗ . is computed as mentioned before. Table 3 illustrates the complexity of these computations.

5 Experiments

We report the perplexity of KN, MKN, SM, and our approach CN using the Finnish (FI), Spanish (ES), German (DE), English (EN), French (FR), portions of the Europarl v7 [Koehn, 2005] corpus, as well as 250MiB, 500MiB, 1,2,4, and 8GiB chunks of English Common Crawl corpus [Buck et al., 2014]. The data was tokenized, sentence split, and the XML markup discarded. As test sets, we used newstest-2014 for all languages except Spanish, for which we used newstest-2013. To avoid the effect of differences in handling Out-of-Vocabulary words in measuring the perplexities, we used a closed vocabulary setup. To measure the KN and MKN perplexities we used the SRILM [Stolcke, 2002] toolkit. And to verify the comparability, we forced the KN assumptions on our model and closely matched (difference ≤ 1) the perplexity numbers reported by SRILM. For benchmarking the memory and time usage of CN against SM, we used the English language datasets varying in size, from 125MiB to 8GiB chunks of Common Crawl. All experiments are done on a single core on Intel Xeon E5-2667 3.2GHz and 180GiB of RAM.

Perplexity As illustrated in Table 4, our approach (CN) consistently outperforms MKN perplexities by a margin of up to 15%. To test against full HPYP inference, we compared against the available implementation of SM. 6 Although CN is initialised by KN and samples from the conditional P(γ+|γ-), it is consistently better than SM. We speculate this is due to poor mixing of SM over the full sampling space involving HPYP tree nodes, and proper mixing of our fast inference method over the smaller sampling space involving only nodes on a single branch γ+. The difference between perplexities across multiple runs of our model were negligible. Comparing SM with KN and MKN reveals a surprising result on some datasets: SM does worse, or is only marginally better regardless of the number of burn-in, or samples. 7

Memory andTime As demonstrated in Figure 4, the memory used by SM is several orders of magnitudes larger than the size of the text and the size of our model in both training and test (query). For instance, on 250MiB dataset CN used (297MiB,108MiB) in training and test, compared with

6 https://github.com/jgasthaus/libPLUMP

7 This doesn't verify the comparison reported in [Wood et al., 2011]. We noticed a critical decision in their experimental setup: while setting a threshold to replace low frequency words with a single token is a popular approach in text processing, in the KN and MKN LM this will cause a range of discount parameters to be zero, eliminating the effect of smoothing and making KN and MKN performworse than their full potential.

(16GiB,16GiB) of SM when it stores only 1 set of HPYP samples. This made it impossible for us to run SM on larger (≥ 2GiB) datasets, or with more samples. In terms of time usage in the training step, our approach is several times faster, i.e., (40×,67×) on (125MiB,250MiB) datasets. Noting that SM tends to get slower on larger datasizes, taking more than 14 days to train on 1GiB dataset, while we only required less than 2 hours. In the test time we are on average around 24× slower noting that SM tends to get slower on larger datasizes, i.e. from 27× faster on 125MiB to 21× on 250MiB. Inference only with 5 samples along each branch and no burnin, affected the perplexities of our approach up to 2% while making the average test speed only 4× slower than SM. Our approach on "load+query" carries roughly a similar pattern excluding the load time. On 1GiB, our query is only 1.8× slower than SM, and we are 2.3× faster on load+query. A significant result which is due to smaller model size, and efficient inference mechanism of our CST-based framework.

6 Conclusion

In this paper we proposed a framework based on compressed suffix trees to represent infinite-order hierarchical Bayesian language models compactly, while developing a fast and memory-efficient approximate inference scheme. Compared with the existing HPYP LMs our approach has several orders of magnitudes lower memory footprint allowing us to apply it on (100×) larger data sizes than the largest data used by HPYP LM. This is achieved by avoiding potential mixing issues, while consistently outperforming the Kneser-Ney family of smoothings by a significant margin.

As our future work, we would like to speedup the inference via approximating Stirling numbers using a separate model to avoid its expensive recursion cost during sampling, as well as exploring continuous space approximations of HPYP.

Acknowledgments

This research was supported by the National ICT Australia (NICTA). The first author would like to thank Wray Buntine for fruitful discussions about sampling in HPYP, and Philip Chan for the support to run the experiments on Monash Advanced Research Computing Hybrid (MonARCH) servers.