Closed-form Baum–Welch for TKF-based evolutionary models

The TKF91 and TKF92 models of sequence evolution combine two continuous-time Markov chains: a linear birth-death-immigration (BDI) process for insertions and deletions, and a finite-state substitution model. Expectation-Maximization (EM) is the standard approach to fitting the substitution part of the model. We report a corresponding closed-form EM algorithm for the indel part: given the expected transition counts of the pairwise hidden Markov model, the optimal insertion rate, deletion rate and fragment-extension probability are functions of a handful of sufficient statistics, with no gradient ascent or line search required. The E-step of this indel-rate EM algorithm reveals a more general result: for any linear BDI process, the bridge statistics (endpoint-conditioned expected counts and dwell times) are closed-form via a score identity on the finite-time transition probabilities. We give a self-contained account of this closed-form Baum-Welch for TKF. We generalize the CherryML composite-likelihood for pairwise alignments to TKF models. This enables TKF-EM time complexity to be independent of training corpus size. We then consider two TKF-based models with extra latent information: MixFrag, a TKF92 with a categorical mixture over geometric fragment lengths (which reduces exactly to TKF92), and its multi-domain generalization MixDom. We develop two trainers for these models: svibw, a stochastic variational Baum-Welch (for both MixFrag and MixDom) that marginalizes alignment and latent structure jointly on minibatches of sequence pairs; and an exact EM (for MixFrag only) that recovers alignment summarization by counting run lengths. On a 20,897-family Pfam training split (1,134,482 pairwise alignments), a two-type MixFrag splits TKF92's single fragment-extension probability into short ($r \sim 0.38$) and long ($r \sim 0.83$) fragment modes and improves held-out likelihood by ~1.9 nats per pair over TKF92 (indel orderings marginalized) at a cost of two parameters, reproducing the familiar empirical observation of the indel length distribution's long tail. On a BAliBASE multiple sequence alignment benchmark, MixFrag and MixDom improve accuracy over TKF92 (by ~2% and ~4\%).