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:
- 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.
- 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.
- 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]].
- 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():
- 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.
- 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.
- 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.
- 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.
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 toRateFree::optimizeWithEM()(ratefree.cpp:323→506). Inside it, per EM step:(
ratefree.cpp:628-658) runs, for every category in every EM step:tree->copyPhyloTree(phylo_tree, true)→tree->initializeAllPartialLh()→ Brent. The scratcharrays' 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 everymodel refresh during tree search. This term scales with taxa and dominates at phylogenomic sizes.
:544computescomputePatternLhCat, then asingle E/M pass). With rates fixed, per-category site likelihoods
L_scare independent of theproportions, so the proportion sub-problem is exactly solvable to convergence on the cached
L_scat zero traversal cost — precisely whatModelMixture::optimizeWeights()(
modelmixture.cpp:4066) already does for mixture weights with its ratio-update trick.optimizeTreeLengthScaling), ~10–30single-category traversals per rate per EM step, with a cold bracket
[MIN_PROP, 1/prop[c]].ncategorysteps (for (step = 0; step < ncategory; ...),:~541) — an arbitrary bound unrelated to convergence (a converged-break already exists).Exact proposed fix
In
RateFree::optimizeWithEM():tree+copyPhyloTree+initializeAllPartialLh()once before the EM loop; inside the per-category loop onlytree->setModel(subst_model),tree->clearAllPartialLH(), refreshptn_freqfrom theposteriors, 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.
_pattern_lh_catbefore touching therates: port the ratio-update loop from
ModelMixture::optimizeWeights()(
modelmixture.cpp:4080-4110) — multiply the cached category likelihoods bynew_prop[c]/prop[c]each sweep and iterate until|Δprop| < 1e-6. Monotone by EM ascent;pure memory passes.
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.
ncategorycap with a fixed generous cap (e.g. 100) and rely on the existingconvergence 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.