Skip to content

Faster free rates #206

Description

@roblanf

Very similar to #205, but this time some speed ups to the free rate optimizer.

Something that's bugged me for a while, but I've lacked the expertise to fix it. I think this should help!

Below written by fable 5.

FreeRate (+R) EM: per-category partial-likelihood reallocation and one-step proportion updates make optimisation much slower than necessary

Version: v3.1.3 (all line refs). Companion issue to the -mwopt/BFGS one — same theme:
sub-problems solvable from cached per-category likelihoods are re-solved with traversals, or here,
with repeated allocations.

Issues

Default optimize_alg_freerate = "2-BFGS,EM" (tools.cpp:7241) routes +R optimisation to
RateFree::optimizeWithEM() (ratefree.cpp:323→506). Inside it, per EM step:

  1. Partial-likelihood arrays are reallocated per category, per EM step. The rate M-step loop
    (ratefree.cpp:628-658) runs, for every category in every EM step:
    tree->copyPhyloTree(phylo_tree, true)tree->initializeAllPartialLh() → Brent. The scratch
    arrays' shape never changes, yet they are allocated and first-touched up to
    ncategory² times per call — multi-GB churn each time at large taxon counts, repeated at every
    model refresh during tree search. This term scales with taxa and dominates at phylogenomic sizes.
  2. Proportions get one EM update per traversal (:544 computes computePatternLhCat, then a
    single E/M pass). With rates fixed, per-category site likelihoods L_sc are independent of the
    proportions, so the proportion sub-problem is exactly solvable to convergence on the cached
    L_sc at zero traversal cost — precisely what ModelMixture::optimizeWeights()
    (modelmixture.cpp:4066) already does for mixture weights with its ratio-update trick.
  3. Each rate is optimised by no-derivative Brent (optimizeTreeLengthScaling), ~10–30
    single-category traversals per rate per EM step, with a cold bracket [MIN_PROP, 1/prop[c]].
  4. The outer EM loop is capped at ncategory steps (for (step = 0; step < ncategory; ...),
    :~541) — an arbitrary bound unrelated to convergence (a converged-break already exists).

Exact proposed fix

In RateFree::optimizeWithEM():

  1. Hoist the scratch tree out of both loops. Build tree + copyPhyloTree +
    initializeAllPartialLh() once before the EM loop; inside the per-category loop only
    tree->setModel(subst_model), tree->clearAllPartialLH(), refresh ptn_freq from the
    posteriors, and rescale. Array shapes are identical across categories (20 states, 1 category),
    including the fused-mixture case, so reuse is safe. Free it once after the loop.
  2. Run the proportion EM to convergence on the cached _pattern_lh_cat before touching the
    rates: port the ratio-update loop from ModelMixture::optimizeWeights()
    (modelmixture.cpp:4080-4110) — multiply the cached category likelihoods by
    new_prop[c]/prop[c] each sweep and iterate until |Δprop| < 1e-6. Monotone by EM ascent;
    pure memory passes.
  3. Warm-bracket the Brent call: the previous rates[c] is a good starting point — bracket
    [rates[c]/3, rates[c]*3] (clamped to the legal range) instead of [MIN_PROP, 1/prop[c]].
    Follow-up (separate PR): replace Brent with a Newton step using the analytic derivative of the
    likelihood w.r.t. a global branch-scaling factor, which the derivative machinery can supply in
    one traversal.
  4. Replace the ncategory cap with a fixed generous cap (e.g. 100) and rely on the existing
    convergence break.

(1) and (2) are behaviour-preserving to numerical tolerance and need no new algorithms;
(3) warm-bracketing likewise. Expected effect: removes the taxon-scaling allocation term entirely
and cuts traversals per EM step several-fold — most visible in tree searches on large matrices,
where +R currently re-pays the full cycle at every model refresh.

Context

Found on a 5,813-taxon × 36,810-site fungal AA analysis; the allocation term in (1) is tens of GB
per category-visit at this scale. It is one reason fixed-R8 pipelines (estimate once, freeze)
currently outperform letting the search keep +R free.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions