Bayesian model merging for stochastic context-free grammar induction

This article describes a weekend project, with a fuller account in the accompanying paper.

Grammar induction requires finding both the productions and their probabilities. Once the productions are fixed, expectation maximisation estimates their probabilities by repeating the inside and outside recursions until the corpus likelihood stops improving. Finding the discrete structure is harder, since the possible sets of productions form a combinatorial space that grows quickly with the number of nonterminals. Stolcke and Omohundro approach this problem in their 1994 paper, Inducing Probabilistic Grammars by Bayesian Model Merging, by beginning with a grammar that memorises the corpus and applying local merge and chunk operations. The implementation here follows their treatment of stochastic context-free grammars in Section 3.3.

Grammar representation and data incorporation

Let a grammar consist of productions with a nonterminal on the left and a sequence of symbols on the right. Each production stores a usage count from which its probability can be obtained. For every observed token $a$, introduce a nonterminal $T_a$ with the single production $T_a \to a$, so that terminals occur only in these lexical productions and every structural right-hand side consists of nonterminals. This follows the representation introduced in Section 3.3 of the paper and allows the two editing operators to act uniformly on structural productions. The OCaml representation is:

type symbol = Term of string | Nonterm of string

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

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

Incorporating the data means folding over the corpus and adding a start production for each distinct sample. Its right-hand side is the corresponding sequence of terminal nonterminals, and its count is the observed frequency of the sample. For $\{ab, aabb, aaabbb\}$, this gives the grammar at the top of Figure 2 in the paper. A one-token sample gives a unit production, which must survive normalisation when it does not belong to 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 models. The paper factors the model prior into a prior over discrete structure and a conditional prior over continuous parameters:

The search chooses a structure by maximising $P(M_S \mid X)$. This posterior is proportional to the structural prior multiplied by the marginal likelihood, with the production probabilities integrated out:

Structural prior

The structural prior penalises the description length of a grammar. The paper specifies a cost of $\log_2 N$ bits for each occurrence of a nonterminal in a right-hand side, where $N$ is the number of nonterminals. The implementation supplies the remaining encoding choices, charging terminal occurrences against the size of the terminal alphabet $\Sigma$:

The first two terms describe the sizes of the nonterminal and terminal alphabets. Within the sum, the encoding records the number of productions for each nonterminal, the length of each right-hand side and the symbols in that sequence, using the alphabet to which each symbol belongs. Converting to nats gives the log prior $-w\,\ell(M_S)\ln 2$, where the weight $w$ defaults to one. Extra productions and longer right-hand sides therefore carry an explicit description cost which favours smaller grammars.

Parameter prior and the marginal likelihood

The parameters are the production probabilities, with an independent symmetric Dirichlet prior for each left-hand side. Integrating a production-count likelihood against this prior gives the following Dirichlet integral for $k$ productions with counts $n_i$:

Here $\Gamma$ is the gamma function, $n = \sum_i n_i$, and $\alpha$ is the concentration assigned to each production, giving $k\alpha$ virtual samples in total. The implementation evaluates $\log\Gamma$ using a Lanczos approximation. It also needs the digamma function below, which it evaluates using an asymptotic expansion and the recurrence $\psi(x+1) = \psi(x) + 1/x$.

Ambiguity and the variational bound

When each corpus string has one parse, the likelihood factors over production counts and this integral gives the exact evidence. For an ambiguous string, the likelihood sums over parses, so there is no single set of counts from which to obtain the evidence by multiplying these per-production terms. The implementation provides two scores for this case. The first computes the exact inside likelihood at fitted parameters and adds a Dirichlet Occam correction, giving an approximation that is exact on unambiguous corpora. The second uses a mean-field variational lower bound:

The variational distribution $q(\theta)$ is Dirichlet with pseudocounts $\beta_i$, and $B = \sum_i \beta_i$. Its geometric mean parameters $\bar\theta_i$ act as unnormalised weights for the inside computation, whose value is $Z$. With the parse distribution chosen consistently with these weights, its entropy combines with the expected log likelihood to give $\log Z$. The bound then subtracts the KL divergence between the variational parameter distribution and the prior:

Each variational iteration computes the geometric mean weights, obtains expected production counts with inside and outside, and adds those counts to the prior pseudocounts. The following sketch shows the updates:

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

The scores have the same form at their respective fixed points, but this does not make them equal in general. At an expectation maximisation fixed point, the first score is the Dirichlet integral of the expected production counts plus the entropy of the posterior over parses. The variational bound has that form at its own fixed point, using counts and a parse posterior obtained from $\bar\theta$ instead of maximum likelihood parameters. They coincide when every parse of a given string has the same production-count vector, including unambiguous corpora and grammars such as $S \to SS \mid a$. On $S \to SS \mid SSS \mid a \mid b \mid ab$, the differences are 0.014 nats at $\alpha = 0.1$ and 1.42 nats at $\alpha = 2$, with opposite signs. The first score is consequently not a bound on the marginal likelihood. An earlier version of this article reported agreement to machine precision on every tested grammar, which described those tests but did not establish agreement in general.

Parsing and parameter estimation

Since right-hand sides are arbitrary sequences rather than Chomsky normal form pairs, parsing uses a generalised inside recursion over spans. Memoisation keys consist of a nonterminal and the two span endpoints, and each production enumerates the ways to divide its span among the symbols on the right. There are no epsilon productions, so every symbol consumes at least one token and a production with several symbols recurses on strictly smaller spans. Unit productions can retain the same span, which is why normalisation must remove unit cycles before parsing. The sketch below uses ordinary probabilities, while the implementation stores inside and outside quantities as logarithms and combines alternatives using a stable log sum:

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 inside value $I(A,i,j)$ sums the probabilities of all parses in which $A$ derives $w_{i+1}\cdots w_j$. Thus $I(S,0,n)$ gives the string probability including ambiguous derivations. A Viterbi variant selects the best parse for inspection, while combining the inside and outside passes gives the expected count of each production:

Here $O(A,i,j)$ is the outside probability for $A$ over the same interval. Expectation maximisation normalises the expected counts separately for each left-hand side, and the resulting production probabilities supply the likelihood and Occam correction.

Operators

The search edits a grammar by choosing either a pair of nonterminals to merge or a sequence of symbols to chunk. Merging replaces both nonterminals with one survivor carrying the union of their productions, which can enlarge the language and shorten its description. If either nonterminal is the start symbol, that name survives. Otherwise the lexicographically smaller name survives, so reversing the input pair gives the same result:

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 introduces a fresh nonterminal for a contiguous sequence of nonterminals. Its production expands to the sequence it replaces, so the operation preserves the generated language, although adding the nonterminal and production usually lowers the score initially. Later merges may recover that cost, requiring the search to look beyond the immediate result. The paper uses beam search for this purpose. It does not specify how to replace repeated occurrences within one right-hand side, so the implementation supports both replacing every nonoverlapping occurrence and replacing only the leftmost occurrence:

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) } ])

Normalisation follows each edit. It combines duplicate productions and sums their counts, removes self loops and unit cycles, and drops nonproductive or unreachable nonterminals. It also removes a production when the left-hand side can derive its right-hand side using the other productions. This last check uses a complete CYK-style decision procedure over sentential forms. It reproduces the deletions in Figure 2 of the paper, including removing $S \to AXB$ once $S \to ASB$ and $X \to AB$ make it derivable.

Search

A chunk may need several subsequent merges before the posterior improves, so accepting only immediate improvements can stop the search before it reaches a useful grammar. The implementation offers the paper's beam search and a best-first frontier. Both retain candidate grammars, score them by the posterior and expand promising candidates, stopping after a fixed number of expansions without improvement. Each proposal is logged with its score and parent, and each accepted step produces Graphviz diagrams of the grammar before and after the edit. In the running example, chunking reduces the score from $-72.35$ to $-74.23$, before merging raises it to $-62.51$. Retaining the intermediate grammar allows the search to reach that improvement.

Experiments

The experiment suite uses the eight test grammars from Cook et al. listed in Table 1 of the paper, including the palindrome language $wcw^R$ with a centre marker. Each is recovered up to a finite comparison length. The suite does not include $ww^R$ without a marker, which Cook et al. report as beyond their method and which Stolcke and Omohundro infer with beam search. It also omits the relative clause grammar at the end of Section 3.3. For the running example, the implementation reaches $S \to T_a T_b \mid T_a S T_b$, generating $a^n b^n$, in two accepted steps. Figure 2 in the paper takes four because normalisation here removes $S \to T_a X T_b$ as soon as the merge makes that production derivable.

Section 4 explains why absolute frequencies affect generalisation. Multiplying all corpus counts by a constant multiplies the likelihood cost of a generalisation by that constant, while leaving its saving in the structural prior fixed. With the three strings in the ratio six to three to one, the implementation generalises to $a^n b^n$ at every tested size up to 200 samples, although the posterior gain decreases as the corpus grows. From 300 samples onwards it accepts no edit and retains the memorising grammar. Counts of $17$, $17$ and $16$ also prevent generalisation, despite totalling the same 50 samples as $30$, $15$ and $5$. In that case the distribution matters: the recursive grammar assigns a geometric distribution to string lengths, fitting decaying frequencies better than nearly uniform ones. The corresponding log likelihoods are $-68.62$ and $-47.74$.

The table gives the results for each target using a fixed seed, beam width three, strings of at most seven symbols and prior concentration 0.1. Its last column records whether the induced grammar is recursive, as required by the nested and embedding targets.

The induced languages match their targets up to the comparison length, and the finite English grammar remains nonrecursive. For $a^{2n}$, $(ab)^n$, $a^n b^n$ and $a^n c\, b^n$, only three or four strings have length at most seven, and all occur among the 50 training samples. Their held-out likelihoods and finite language comparisons therefore establish fit within that range; the induced recursion supplies the evidence of generalisation. The arithmetic target has another limitation: the search merges $+$ and $*$ into one operator class and collapses the three precedence levels into a single nonterminal. This produces the right strings with an ambiguous grammar that represents no precedence. The runs are deterministic given the seed and beam width.

Reproduction notes

The paper leaves some choices unspecified. The implementation supplies a description-length encoding around the stated $\log N$ cost per nonterminal occurrence and exposes a prior weight because the paper writes $\exp(-\ell)$ while measuring $\ell$ in bits. It also exposes the unspecified Dirichlet concentration and supports both interpretations of chunking repeated occurrences through a flag. The accompanying report records these choices and the complete derivability check used in place of the earlier bounded pruning search.