Inducing stochastic context-free grammars by Bayesian model merging

Based on a weekend project, which you can see here.

Grammar induction splits into two problems, and the harder one is finding the discrete structure. Parameter estimation is comparatively settled because once the productions are known, expectation maximisation fits their probabilities by iterating the inside and outside recursions until the corpus likelihood stops improving. Choosing the productions is a search over a combinatorial space. The 1994 paper by Stolcke and Omohundro, Inducing Probabilistic Grammars by Bayesian Model Merging, turns that search into a sequence of local decisions by starting from a grammar that memorises the corpus and then merging or chunking its parts. The implementation described here follows the stochastic context-free grammar treatment in Section 3.3 of the paper.

Grammar representation and data incorporation

The grammar is a set of productions whose left-hand sides are nonterminals and whose right-hand sides are sequences of symbols. Every production carries a usage count from which a probability can be derived. No terminal ever appears inside a structural right-hand side, so every observed token \(a\) receives a nonterminal \(T_a\) with the single production \(T_a \to a\), which leaves all structural productions with right-hand sides built only from nonterminals. This is the choice the paper makes at the start of Section 3.3, and it is what lets the same two operators apply uniformly later. The representation in OCaml is small enough to quote in full:

type symbol = Term of string | Nonterm of string

type production = {
  lhs : string;
  rhs : symbol list;
  count : float;
}

type t = { start : string; productions : production list }

Data incorporation is then a fold over the corpus, in which each distinct sample becomes one production from the start symbol whose right-hand side is the sequence of terminal nonterminals, and the observed sample count is stored on that production. The corpus \(\{ab, aabb, aaabbb\}\) produces exactly the grammar at the top of Figure 2 in the paper. A sample consisting of a single token produces a unit production, which is why the normalisation step described below keeps unit productions that do not lie on a cycle.

let initial_grammar ?(start = "S") samples =
  let term_nts = Hashtbl.create 64 in
  let prods = ref [] in
  List.iter
    (fun (tokens, count) ->
      let rhs =
        List.map
          (fun tok ->
            match Hashtbl.find_opt term_nts tok with
            | Some nt -> Nonterm nt
            | None ->
                let nt = "T_" ^ tok in
                Hashtbl.replace term_nts tok nt;
                prods := { lhs = nt; rhs = [ Term tok ]; count = 1.0 } :: !prods;
                Nonterm nt)
          tokens
      in
      prods := { lhs = start; rhs; count } :: !prods)
    samples;
  make ~start (List.rev !prods)

The Bayesian score

Bayes' rule gives the posterior over structures, and the paper separates the prior into a term over the discrete structure and a term over the continuous parameters.

The search maximises the structure posterior \(P(M_S \mid X)\), which is proportional to the product of the structural prior and the marginal likelihood obtained by integrating out the parameters.

Structural prior

The structural prior is a description length, and the paper states that every occurrence of a nonterminal in a right-hand side costs \(\log_2 N\) bits, where \(N\) is the number of nonterminals. The implementation adds the remaining encoding choices explicitly.

The first two terms encode the sizes of the nonterminal and terminal alphabets, the third encodes the number of productions of each nonterminal, the fourth encodes the length of each right-hand side, and the final sum charges for the symbols themselves. The prior is then \(P(M_S) = \exp(-\ell(M_S)\ln 2)\), so a grammar with more productions or longer right-hand sides pays for each extra choice it represents. This is the Occam factor made concrete, and it is the only pressure that pushes the search towards smaller grammars.

Parameter prior and the marginal likelihood

The parameters are the production probabilities, and their prior is a product of symmetric Dirichlet distributions, one for each left-hand side. Integrating the multinomial likelihood against that prior yields the Dirichlet multinomial, which for a left-hand side with \(k\) productions and counts \(n_i\) is

where \(\Gamma\) is the gamma function and \(\alpha\) is the concentration of the prior, which plays the role of a number of virtual samples spread evenly over the productions. The implementation uses a Lanczos approximation for \(\log\Gamma\) and an asymptotic expansion with the recurrence \(\psi(x+1) = \psi(x) + 1/x\) for the digamma function that appears below.

Ambiguity and the variational bound

That closed form is exact when each string has a single parse because the corpus likelihood then factors into a product over productions. Ambiguous strings break the factorisation since the likelihood is a sum over parses and the integral of a sum does not decompose into the product of the per-production integrals. Rather than pretend otherwise, the implementation offers two scores. The first computes the exact inside likelihood at the fitted parameters and adds the Dirichlet Occam correction, which is exact for unambiguous corpora. The second makes the approximation explicit as a mean-field variational lower bound.

Here \(q(\theta)\) is the Dirichlet posterior with pseudo-counts \(\beta_i\), the geometric means \(\bar\theta_i\) are the variational parameters, \(B = \sum_i \beta_i\), and \(Z\) is the inside value computed under those unnormalised weights. At the fixed point the entropy of the parse distribution and the expected log likelihood cancel, which leaves the bound equal to \(\log Z\) minus the divergence between the posterior and the prior,

The variational loop alternates between computing the geometric-mean parameters, running inside and outside to get expected counts, and updating the pseudo-counts, which is short enough to state as a sketch:

for _ = 1 to em_iters do
  set_theta_bar ();
  let acc = expected_counts corpus in
  for i = 0 to np - 1 do beta.(i) <- alpha +. acc.(i) done;
  kl := dirichlet_kl beta
done;
elbo := log_z theta_bar () -. !kl

On every grammar, corpus and concentration tested, the two scores agree to machine precision, and both reduce to the Dirichlet multinomial when the corpus is unambiguous. That agreement is reassuring because it means the simpler score was already a bound rather than an ad hoc correction, and the variational mode simply makes the bound explicit.

Parsing and parameter estimation

Right-hand sides are sequences rather than Chomsky normal form pairs, so the parser is a generalised inside recursion over spans instead of a binary CYK loop. The recursion memoises on the nonterminal and the two endpoints, and for each production it enumerates the ways to divide the span among the symbols. Every symbol consumes at least one token because there are no epsilon productions, which bounds the split enumeration and also prevents a recursive nonterminal from calling itself on the same span:

let rec inside_nt nt i j =
  if j <= i then 0.0
  else match Hashtbl.find_opt imemo (nt, i, j) with
    | Some v -> v
    | None ->
        let v = List.fold_left (fun acc pi ->
          acc +. probs.(pi) *. seq_inside prods.(pi).rhs i j) 0.0 (by_lhs nt) in
        Hashtbl.replace imemo (nt, i, j) v; v

and seq_inside rhs i j =
  match rhs with
  | [] -> if i = j then 1.0 else 0.0
  | Term a :: rest ->
      if i < n && String.equal w.(i) a then seq_inside rest (i + 1) j else 0.0
  | Nonterm x :: rest ->
      let acc = ref 0.0 in
      for k = i + 1 to j - List.length rest do
        acc := !acc +. inside_nt x i k *. seq_inside rest k j
      done;
      !acc

The value of \(I(A,i,j)\) is the probability that \(A\) derives the span \(w_{i+1}\cdots w_j\) summed over all parses, so the top-level call \(I(S,0,n)\) is the string probability and it accounts for ambiguity rather than selecting a single derivation. The same table answers three questions, with the inside value giving the likelihood, the Viterbi variant returning the single best parse for inspection, and the inside and outside passes together giving the expected count of each production,

where \(O(A,i,j)\) is the outside probability of \(A\) spanning the same interval. Expectation maximisation renormalises those counts per left-hand side, and the resulting probabilities feed both the likelihood and the Occam correction.

Operators

Two operators generalise a grammar, and both are local edits. Merging two nonterminals replaces them with a single symbol that inherits the union of their productions, which can enlarge the language while shortening the description. The survivor is the start symbol when one of the pair is the start, and otherwise the lexicographically smaller name, so the result does not depend on the order in which the pair was generated:

let apply_merge g a b =
  let survivor, removed =
    if String.equal a g.start then (a, b)
    else if String.equal b g.start then (b, a)
    else if String.compare a b <= 0 then (a, b) else (b, a)
  in
  let map = function
    | Nonterm x when String.equal x removed -> Nonterm survivor
    | s -> s
  in
  let prods = List.map (fun p ->
    let lhs = if String.equal p.lhs removed then survivor else p.lhs in
    { p with lhs; rhs = List.map map p.rhs }) g.productions
  in
  copy_with g prods

Chunking abbreviates a contiguous sequence of nonterminals with a fresh nonterminal. It preserves the generated language exactly, since the new production expands back to the sequence it replaced, but it usually lowers the score on its own by adding a production and a nonterminal. That local dip is precisely why the search has to look ahead, and it is the reason the paper reaches for beam search rather than greedy merging. The implementation replaces either every non-overlapping occurrence of the sequence or only the leftmost one, because the paper does not say what to do when a sequence repeats inside a single right-hand side:

let apply_chunk ?(occurrence = All_occurrences) g seq =
  let name = fresh_name g "X" in
  let replace = match occurrence with
    | All_occurrences -> replace_all
    | First_occurrence -> replace_first
  in
  let total = ref 0 in
  let prods = List.map (fun p ->
    let rhs, n = replace seq name p.rhs in
    total := !total + n; { p with rhs }) g.productions
  in
  copy_with g (prods @ [ { lhs = name; rhs = seq;
                           count = float_of_int (max 1 !total) } ])

After either operator the grammar is normalised, and duplicate productions are merged and their counts summed, self loops and unit cycles are removed, nonproductive and unreachable nonterminals are dropped, and a production is removed when its left-hand side can derive its right-hand side without it. That last check is a complete CYK-style decision procedure over sentential forms rather than a bounded search, so it removes exactly the productions that the paper's Figure 2 drops after its merges, including the production \(S \to AXB\) that disappears once \(S \to ASB\) and \(X \to AB\) are present.

Search

Greedy merging is not enough for context-free grammars because a chunking step typically needs several following merges before the posterior improves. The paper therefore uses beam search, and the implementation offers that alongside a best-first frontier. Both keep a set of candidate grammars, score each one by the posterior, expand the most promising candidates, and stop after a fixed number of expansions without an improvement. Every proposal is logged with its score and its parent, and each accepted step is written out as a pair of before and after Graphviz diagrams. In the running example the first chunk lowers the score from \(-72.35\) to \(-74.23\), and the merge that follows raises it to \(-62.51\), which is the local dip the beam exists to survive.

Experiments

The paper's Table 1 lists eight target grammars, and the implementation recovers each of them up to a comparison length, including the palindrome language \(wcw^R\) that Cook et al. considered out of reach for their method. The recursive relative-clause grammar at the end of Section 3.3 is recovered as well. The running example converges to \(S \to T_a T_b \mid T_a S T_b\), which is the language \(a^n b^n\), in two accepted steps. Section 4 of the paper argues that absolute frequencies, and not only relative frequencies, control generalisation. The implementation reproduces that effect exactly, and the three strings with counts \(30\), \(15\) and \(5\) induce the recursive grammar, while near uniform counts of \(17\), \(17\) and \(16\) block the leap, because the same relative distribution with a larger absolute weight makes each generalisation cost more likelihood.

The table records the outcome for every target, using a fixed seed, a beam of four, and a prior concentration of 0.1. The final column reports whether the induced grammar is recursive, which is the property the nested and embedding targets are there to test.

Every target is recovered exactly up to the comparison length. The held-out average log-likelihood per token is negative but small for the recursive languages, and the finite English grammar is correctly not recursive. The full set of runs is deterministic given the seed and the beam width.

Reproduction notes

Three details are left open by the paper, and the implementation states its choices rather than hiding them. The description-length encoding is one reasonable instantiation, since only the \(\log N\) cost per nonterminal occurrence is given, and the prior weight is exposed as a parameter because the paper writes the prior as \(\exp(-\ell)\) with \(\ell\) measured in bits. The concentration of the Dirichlet prior is not specified, so it is also a parameter. The multi-occurrence behaviour of chunking is not specified, so both readings are available behind a flag. Each of these choices is recorded in the report that accompanies the code, together with the observation that the two likelihood scores coincide to machine precision and the complete derivability check that replaced the bounded pruning search.