Concurrent (single-solve) fitting for multidms

[jaredgalloway]

Concurrent (single-solve) fitting for multidms

Status

READY

Goal

The multidms fitting procedure (multidms.jaxmodels.fit) currently optimizes
the model via block coordinate descent: each outer iteration cycles
sequentially through four blocks — calibration (α, logθ), β0, β
non-bundle, and β bundle — repeating for up to block_iters outer iterations.

We want to test an alternative that optimizes the entire model
concurrently in a single proximal-gradient solve, and compare it head-to-head
against block CD on quality, sparsity, and speed. This is a methods
experiment: the comparison is the deliverable, not a silent replacement.

Background: the current architecture

fit() (jaxmodels.py:364) runs, per outer iteration k in range(block_iters):

1. Calibration block — α (GE scaling) + logθ (count overdispersion),
   jaxopt.GradientDescent on objective_part.
2. β0 block — per-condition intercepts, GradientDescent, optional ridge.
3. β non-bundle block — mutation effects at sites that are WT in all
   conditions, jaxopt.ProximalGradient with the fusion-lasso prox.
4. β bundle block — mutation effects at sites non-WT in some condition,
   same ProximalGradient / same prox.

Convergence is declared when the relative change in objective_total drops
below block_tol.

The fusion penalty (fusionreg, an L1 on the shift β_d − β_ref between
each non-reference condition and the reference) is non-smooth and lives
entirely in the proximal operator prox_block (jaxmodels.py:525), which
runs only inside the two β blocks. The smooth penalties (l2reg on β,
beta0_ridge on β0 differences) are folded into the objectives.

Dynamic scale. Each outer iteration recomputes scale = |objective_total|
so the lasso threshold fusionreg / scale stays calibrated as the objective
shrinks (jaxmodels.py:631-643).

The concurrent approach

Replace the 4-block × block_iters nested loop with one
jaxopt.ProximalGradient.run call over the full parameter pytree.

Architecture

fit_concurrent(data_sets, ...)   # SAME signature/returns as fit()

  1. Warmstart (Ridge)  ->  initial (β0, β, α, logθ)   [identical to fit()]
  2. Reparameterize     ->  free vars = (β_ref, {Δ_d}, β0, α, logθ)
  3. scale = |objective_total(init_model)|             [computed ONCE,
              with the same > 1e-30 guard as fit() (jaxmodels.py:643):
              scale = raw_obj if raw_obj > 1e-30 else 1.0]
  4. ONE jaxopt.ProximalGradient.run(full_pytree):
       fun  = smooth objective: loss + l2·Σ‖β_d‖² + β0_ridge·Σ(β0_d−β0_ref)²
              (NO fusion term — it lives in the prox)
       prox = prox_lasso on Δ_d leaves (threshold fusionreg/scale),
              identity (prox_none) on β_ref, β0, α, logθ;
              optional β box-clip via Δ adjustment
  5. Reconstruct β_d = β_ref + Δ_d  ->  return standard Model
  6. Build trajectory DataFrame (same schema as fit())

No outer loop. jaxopt's FISTA iterations (acceleration=True default) and
backtracking line search replace the manual block sweeps; maxiter replaces
block_iters as the iteration budget.

Three load-bearing decisions

1. Fuse the bundle/non-bundle β split into one prox.
The split is purely computational — both blocks apply the identical
prox_block on disjoint index subsets, with no mathematical coupling
(confirmed: bundle_idxs at jaxmodels.py:690 only partitions which sites ever
appear as variants). It existed so block CD could run the non-bundle sites —
which never touch the reference's fusion term — as a faster independent
sub-problem. In a single concurrent solve there is no per-block sub-problem to
accelerate, so the split is dead structure. Apply prox_lasso to all
Δ-shifts at once. The union of the bundle and non-bundle index sets is the full
β vector, so thresholding the whole Δ_d leaf is exactly equivalent to
thresholding the two slices separately (soft-threshold leaves zeros at zero —
non-bundle sites where Δ_d ≡ 0 stay zero).

2. Fix the prox scale once, from the warmstarted initial objective.
The dynamic scale refresh existed to recalibrate the threshold across outer
iterations; with one solve there are no outer iterations to recalibrate across.
Compute scale = |objective_total(initial warmstarted model)| once and hold it
fixed through the solve. This preserves calibration compatibility with all
existing fusionreg hyperparameter values (dropping scaling entirely would
force recalibration of the whole spike-pipeline grid) and keeps the solve
genuinely single-pass (a refresh loop would reintroduce the alternation we are
removing). The backtracking line search adapts the step size independently, so a
fixed threshold denominator is sound. Diagnose via the trajectory; add
refreshing only if convergence data demands it.

3. Reparameterize the free variables to (β_ref, {Δ_d}) — the correctness
linchpin.
Soft-thresholding Δ = β_d − β_ref while β_ref is a free, simultaneously
gradient-updated variable is not the prox of the joint fusion penalty
Σ_d ‖β_d − β_ref‖₁ in (β_ref, β_d) coordinates — that penalty is not
separable in those leaves (β_ref appears in every term). Block CD sidesteps this
because each β block freezes β_ref as a constant during its prox step. A naive
concurrent port would be mathematically wrong.

The standard fix: make the solver's free variables (β_ref, {Δ_d}) and have the
forward pass reconstruct β_d = β_ref + Δ_d. Then the fusion penalty is plain
lasso on the separable Δ_d leaves — prox_lasso is its exact closed-form
prox — and identity on β_ref. This makes the concurrent prox-gradient step a
genuine composite-objective solve, not an approximation of block CD.

Reparameterization mechanics

The reparameterization lives at the solve boundary, internal to
fit_concurrent. The public Model is unchanged, so ModelCollection,
plot.py, and all downstream code are untouched — only the optimization
variables are (β_ref, Δ_d).

free pytree  P = {
  β_ref : Float[n_mut]          # reference condition's β
  Δ     : {d: Float[n_mut]}     # shift for each NON-reference condition
  β0    : {d: Float[]}          # all conditions (incl. ref)
  α     : Float[] or {d: ...}   # shared (default) or per-condition
  logθ  : {d: Float[]}
}

to_model(P) -> Model:           # reconstruct a standard Model for fun/loss
  φ[ref]   = Latent(β0[ref], β_ref)
  φ[d≠ref] = Latent(β0[d],   β_ref + Δ[d])
  Model(φ, α, logθ, reference_condition, global_epistasis)

fun(P)  = loss(to_model(P))
          + l2reg · Σ_d ‖β_d‖²              # on RECONSTRUCTED β_d, not Δ_d
          + beta0_ridge · Σ_d (β0_d − β0_ref)²
          #  ^ smooth only; NO fusion term

prox(P, hp, η):                # jaxopt calls prox(params, hyperparams, scaling)
  # jaxopt applies ONE prox to the WHOLE pytree; there is no built-in
  # "prox per named leaf". So this is a CUSTOM structured prox that walks
  # P's named sub-trees explicitly: prox_lasso on the Δ sub-tree only,
  # identity (prox_none) on β_ref, β0, α, logθ.
  Δ[d]  <- prox_lasso(Δ[d], fw[d]·fusionreg/scale, η)   # only nonsmooth op
  if beta_clip_range:                                    # usually None
    Δ[d] <- clip(β_ref + Δ[d], range) − β_ref            # clip on reconstructed β_d
    β_ref <- clip(β_ref, range)
  # β0, α, logθ, β_ref: identity
  return P

The hyperparameters_prox dict threads the same fields the current
prox_block reads (jaxmodels.py:525-530): fusionreg, scale,
fusion_weights, beta_clip_range. The difference is purely structural — the
prox now dispatches by pytree sub-tree (Δ vs everything else) instead of by
β-index slice (bundle vs non-bundle).

Correctness details:

• fun carries no fusion term. Fusion L1 lives entirely in the prox — the
  textbook composite split min f(x) + g(x) with f = fun, g = fusion
  lasso realized by the prox. The existing objective_total (which does
  include fusion) stays as-is, used only for the one-time scale computation
  and trajectory logging — never as the solver's fun.
• l2 on reconstructed β_d, not Δ_d. Matches current objective_block
  semantics (jaxmodels.py:498) so l2 calibration is identical to block CD.
• β box-clip. Current code clips β_d directly; under (β_ref, Δ_d) that is
  not separable, so clip β_ref and clip the reconstructed β_d by adjusting
  Δ_d. beta_clip_range defaults to None, so this path is usually inert.
• Warmstart → free vars. Warmstart yields per-condition (β0_d, β_d) exactly
  as today; set β_ref = β_warmstart[ref], Δ_d = β_warmstart[d] − β_ref. The
  two fitters thus start from an identical model — any difference is purely
  the optimization path.

Solver construction

opt = jaxopt.ProximalGradient(
    fun=fun,                # smooth: loss + l2 + β0_ridge (NO fusion)
    prox=prox,              # lasso on Δ leaves, identity elsewhere
    maxiter=maxiter,        # replaces block_iters
    tol=tol,                # replaces block_tol; state.error is jaxopt's
                            # gradient-mapping / fixed-point residual.
                            # Confirm the exact normalization against the
                            # INSTALLED jaxopt version rather than assuming.
    stepsize=0.0,           # 0 => backtracking line search (default)
    acceleration=True,      # FISTA (default); ISTA fallback if GE-Sigmoid
                            # nonconvexity causes oscillation
    # plus maxls, decrease_factor surfaced from solver_kwargs
)
P_final, state = opt.run(P_init, hyperparameters_prox)

• block_iters/block_tol → maxiter/tol. The cal_kwargs/ge_kwargs split
  collapses to a single solver_kwargs dict — there is one solver now, not four.
  This is the one intentional signature difference from fit(); the
  signature is otherwise identical (same data/regularization/init/GE arguments).
• acceleration=True is exposed because the GE Sigmoid makes fun
  nonconvex — FISTA's O(1/k²) guarantee does not formally hold there. The
  trajectory reveals oscillation; ISTA (acceleration=False) is the documented
  fallback.

Output schema

fit() logs one row per outer block-iteration with four block-specific
diagnostic groups (calibration_error, beta0_error, beta_nonbundle_error,
beta_bundle_error, each with *_stepsize/*_iter_num). A single solve has no
blocks.

Decision: emit the same columns for drop-in ModelCollection
compatibility, populating the four now-defunct block diagnostic groups with the
single solver's state values (same value across the columns) and documenting
this in the docstring. Log one row per solver iteration so iteration,
objective_total_trajectory, loss_*, sparsity_*, and per-condition
parameter columns remain populated — plot.py convergence plots keep working
unchanged.

Per-iteration logging mechanism (resolve at implementation time).
jaxopt.ProximalGradient.run returns only the final (params, state) — it does
not expose per-iteration history by default. Two viable mechanisms; pick one and
note it in the docstring:

• Manual opt.update loop — call opt.init_state then loop opt.update
  for maxiter steps, recording a trajectory row each step. This is a logging
  loop, not block alternation: it is still one concurrent prox-gradient solve
  over the full pytree (the "no outer block loop" property holds — the
  reparameterization and single shared step size are unchanged). Preferred,
  because it makes the same rich per-iteration logging fit() produces
  straightforward.
• jaxopt callback — pass a callback that appends rows; lighter but JAX
  side-effect/tracing constraints make per-iteration Python-side logging
  awkward. Fallback only.

The manual-update loop is the recommended choice and does not reintroduce
block coordinate descent — the distinction is "one solver stepped N times with
logging" vs "four solvers alternated N times."

Integration

In scope (this issue): the library function + its tests. Add
jaxmodels.fit_concurrent() returning the same (Model, trajectory DataFrame) schema as jaxmodels.fit(), living next to fit() in
multidms/jaxmodels.py. fit() (block CD) is left unchanged so the
head-to-head baseline is preserved. The validation tests below run directly
against fit_concurrent / fit at the jaxmodels layer.

Out of scope (documented follow-up): spike-pipeline strategy: "concurrent"
wiring. Selecting the concurrent fitter from the spike pipeline is NOT a peer
toggle to independent / continuation and must not be presented as one. Those
two strategies are selected at the notebook level — the Snakefile picks
fit_models.ipynb vs fit_models_path.ipynb, which call two different
model_collection functions — and both ultimately call the same
jaxmodels.fit, because multidms.Model.fit (model.py:240) hardcodes it.
Adding jaxmodels.fit_concurrent does not make it reachable from the pipeline.
Full wiring is a separate follow-up issue requiring four plumbing layers:

1. a strategy/selector param on multidms.Model.fit to dispatch
   jaxmodels.fit vs jaxmodels.fit_concurrent (model.py:240 is currently
   hardcoded);
2. threading that param through model_collection.fit_models / _fit_fun
   (model_collection.py:159);
3. a third branch in the Snakefile strategy validation and FIT_NOTEBOOK
   selection (currently a binary continuation else fit_models.ipynb);
4. config-schema handling for the cal_kwargs/ge_kwargs → solver_kwargs
   collapse (the pipeline currently populates both legacy dicts via
   build_fit_params), either by adding a solver_kwargs config key or having
   fit_concurrent accept and merge/ignore the two legacy dicts.

The A/B comparison in this issue is therefore run at the library/test layer
(call both fitters on the same data_sets in a test or a small experiment
script), not via the spike-pipeline config. Pipeline integration follows once
the method is validated.

Validation plan

The comparison is the deliverable. All checks run at the jaxmodels layer
(call fit and fit_concurrent directly), per the Integration scope above.

(a) pytest correctness tests

These certify the math, not just smoke-test. Use float64 (the module enables it,
jaxmodels.py:29) and a tiny hand-checkable fixture (≈2 conditions, ~4
mutations, ~6 variants) so active sets can be verified by hand.

• Reparameterization round-trip. to_model(decompose(warmstart)) == warmstart model — β reconstruction β_d = β_ref + Δ_d is exact. Cheap guard
  on the coordinate change itself.
• Finite-difference gradient check on fun, GE-Sigmoid ON (T2). On the
  tiny fixture with global_epistasis=Sigmoid() and random P,
  jax.grad(fun)(P) must agree with a central finite-difference estimate to
  ~1e-5. Isolates the reconstruction and the "l2 on reconstructed β_d, not Δ_d"
  / β0-ridge terms from the optimizer — a sign/index error in to_model shows
  up here, not in any end-to-end test.
• KKT / gradient-mapping stationarity at the returned solution (T1 — the
  most important addition). For a converged (β_ref, {Δ_d}, β0, α, logθ),
  compute the prox-gradient fixed-point residual
  R = P − prox(P − η·∇fun(P), η) at the solver's final η; require
  ‖R‖∞ < 1e-4. Equivalently, check the subgradient KKT condition per Δ leaf:
  for Δ_di ≠ 0, ∇fun_i = −fw[d]·fusionreg/scale·sign(Δ_di); for Δ_di = 0,
  |∇fun_i| ≤ fw[d]·fusionreg/scale. This is the only test that certifies
  the returned point is stationary for the composite objective independent of
  block CD — a "wrong prox" can still converge to a non-stationary point.
• fusionreg=0 equivalence to block CD. With no fusion the prox is inert and
  the problem is smooth; concurrent and block CD must reach the same optimum
  within tol. (Note: cannot catch a fusion-prox bug — that's what the next test
  is for.)
• Nonzero-fusionreg equivalence to block CD (T3). On the tiny fixture
  with Identity() GE (convex → single optimum, no nonconvex ambiguity), pick a
  fusionreg that thresholds a known subset of shifts to exactly zero (above
  one Δ's magnitude, below another's). Run both fit and fit_concurrent to
  tight tol with acceleration=False. Require: matching support set AND
  ‖β_concurrent − β_blockCD‖∞ < tol AND |Δobjective_total| < tol. This is the
  test that actually exercises the reparameterized lasso against the trusted
  baseline.
• ISTA monotonic objective decrease (T4). Tiny fixture, GE-Sigmoid,
  acceleration=False; from the logged trajectory, objective_total must be
  non-increasing across iterations (ISTA + backtracking guarantees descent even
  on the nonconvex GE fun; FISTA does not). A violation flags a wrong prox, a
  mis-scaled fusionreg/scale threshold, or line-search misconfiguration — and
  is the natural smoke test for the "ISTA fallback if oscillation" decision.
• Prox idempotency (cheap). prox(prox(P)) == prox(P) on the Δ leaves —
  confirms the soft-threshold is a true prox, not e.g. applied with a doubled
  threshold.

(b) library-layer A/B comparison (the research question)

Run on realistic data (e.g. an existing spike dataset loaded directly, not via
the pipeline config), comparing fit vs fit_concurrent:

Check	What it proves
Identical warmstart, identical inputs	Divergence is pure optimization path, not init.
Sparsity-pattern agreement at matched fusionreg	Concurrent recovers comparable lasso sparsity; large disagreement flags a prox/scale bug.
Objective + per-condition loss head-to-head	Concurrent should reach ≤ block-CD objective (or within noise) to qualify as "as good."
Wall-clock + iteration count	The research question: is one FISTA solve faster than block_iters × 4-block sweeps to equal quality?
Data-poor condition behavior	The motivating pathology (strong shift lasso distorting data-poor conditions); compare shift estimates.

Scope / non-goals

• No parameter rescaling / preconditioning in v1. The shared single step size
  over heterogeneous-scale leaves (scalar α vs thousands of β) is a known
  proximal-gradient failure mode (the line search can collapse to the steepest
  block's Lipschitz scale, starving α/logθ). Per the "start simple,
  diagnose later" decision: ship the plain shared-step solver with backtracking,
  measure convergence via the trajectory, and add scaling or a retained
  separate calibration optimizer only if diagnostics show a problem.
• No removal or modification of fit() / block CD. It is the baseline.
• No smooth-L1 surrogate. We keep the exact lasso via the proximal operator.
• No change to the public Model / ModelCollection / plot.py API.
• No spike-pipeline wiring in this issue. Model.fit dispatch,
  model_collection threading, the Snakefile strategy branch, and the
  solver_kwargs config-schema change are a documented follow-up (see
  Integration). This issue delivers jaxmodels.fit_concurrent + tests; the A/B
  comparison runs at the library layer.

Key references (jaxmodels.py)

• fit() block loop: 364, 624 (outer loop), 646-757 (the four blocks)
• prox_block (fusion lasso prox): 525-544
• objective_total (incl. fusion; kept for scale + logging): 502-522
• objective_block (l2 on reconstructed β): 487-500
• bundle_idxs computation: 690-693
• dynamic scale recompute: 631-643
• Latent (β0, β; forward pass): 98-189
• Model (φ, α, logθ; predict_score): 225-258
• loss functions (functional_score_loss, count_loss): 297-361
• Latent.warmstart (Ridge init): 145-173