Skip to content

Model fitting speed ups - #140

Merged
James-Thorson merged 9 commits into
devfrom
dev-speedup
Oct 7, 2026
Merged

James-Thorson merged 9 commits into
devfrom
dev-speedup

Conversation

@seananderson

Copy link
Copy Markdown
Collaborator

I had Claude's Opus help with some pretty extensive benchmarking and profiling, trying to figure out why tinyVAST was slower than sdmTMB on SPDE models once data and mesh were both reasonably sized. Overall, these small changes result in the fitting being about ~25-30% faster than the current dev/main branches with somewhat reduced memory footprint too. It's now much closer to sdmTMB, but still a bit slower. The fn() and gr() call timings at least basically match now. More notes below.

Here's what this pull request does:

  • Vectorizes apply() over the A-matrix on the R side (easy win)
  • Moves tau scaling into the GMRF densities, which it turns out is faster
  • Makes DATA_UPDATE(weights_i) optional; turns out it's slower and I assume it's not usually needed
  • Change nll observation error lines from -1 * density to -density. This looks completely trivial but strangely seems to interact with DATA_UPDATE for an additional speed up?? I don't know. Maybe it didn't do anything, but it doesn't hurt at least.

If you want me to revert the DATA_UPDATE change, we can pull that out of this PR. The others are more benign.

Compared to sdmTMB on some spatial/spatiotemporal models, it still seems to often take more nlminb iterations to converge, even with equivalent starting values. It could be due to the parameterization of marginal SD as leading parameter rather than 'tau', but that's a more major structural change that may not be worth it. I haven't touched that.

I also fixed two unit tests that were already erroring on my machine.

- With help from Claude's Opus 5.5, did some extensive benchmarking
  trying to figure out why tinyVAST was slower than sdmTMB for
  relatively large SPDE models. Here's what we came up with.
- Overall, ~25-30% faster
- Somewhat reduced memory too

Changes:
- Vectorize apply() over A-matrix on R side (easy win)
- Put tau in GMRF densities
- Make DATA_UPDATE(weights_i) optional; it's slower
- Change nll lines from -1 * d() to -d();
  trivial but strangely seems to interact with DATA_UPDATE speed-up
- Models also seem to often take more nlminb iterations to converge;
  Could be due to parameterization of marginal SD as
  leading parameter rather than 'tau', but that's a more major structural
  change that may not be worth it.
Future Rho/Gamma and Q_ss were built from the fitted parameters
rather than the sampled draw, so parameter uncertainty didn't
propagate into projected process errors
@seananderson

Copy link
Copy Markdown
Collaborator Author

I added some major speedups to project() in my last push to this branch. Here's my benchmark:

project() (s)
sdmTMB local 1.04
tinyVAST main 132.9
tinyVAST dev-speedup 1.13

which works out to a ~118x speedup!

This additional work is Claude assisted (so do check carefully), but it looks right to me and in my tests the returned values were the same where comparable and the overall distributions are the same. I should set up some long project comparisons between sdmTMB and tinyVAST as a final check.

@Arnaud-Gruss tagging you here since we were discussing it in another thread.

@James-Thorson
James-Thorson merged commit bffa0ab into dev Oct 7, 2026
5 checks passed
@James-Thorson

Copy link
Copy Markdown
Collaborator

I'll admit that I don't fully follow the all of the changes in ´project´. @Arnaud-Gruss, could you please carefully check the output of project?

@seananderson besides the changes in the linear algebra being used in ´project´, the other changes all make sense (even if I don't necessarily understand why they speed it up). Thanks for doing the digging! I should have been more cautious about adding the DATA_UPDATE option ... its there to allow testing for stochastic gradient descent for large data sets, but prolly should have stayed as an exploration.

@seananderson

Copy link
Copy Markdown
Collaborator Author

Here's Claude's summary for future reference:

1. 3a1235e: Add nsim and build TMB objects once (main speedup)

  • Before: each sample needed a separate project() call. Each call ran tinyVAST() twice (once to get the padded-time structure, once with the merged parameters) and then a full predict(), which builds yet another TMB object. So each sample paid for about three model builds and tapes.
  • After: all nsim parameter draws come from a single MC(n = nsim) or rmvnorm_prec(n = nsim) call and are stored as the columns of parmat. The padded model and a prediction object (predobj) are each built once. The per-sample loop does only lightweight R work plus $report() calls.
  • Cost goes from roughly nsim × (model builds + predict) to one build plus nsim cheap evaluations.
  • The examples now use project(..., nsim = n_sims) instead of a loop. With nsim > 1, the return value is a matrix with one column per sample.

2. 8cdd288: Use the Kronecker structure when conditioning future epsilon

  • Before: the full space-time precision Q_hh = kronecker(Q_kk, Q_ss), of size (S·T·C)², was formed and passed to conditional_gmrf(). That requires a large sparse factorization in every sample.
  • After: every site is observed at each past time, so Q_ss cancels out of the conditional mean, which simplifies to -(Q_uu⁻¹ Q_uo ⊗ I) x_obs. The mean is then a solve against the small time × variable matrix Q_kk, applied to all sites at once. For simulation, columns are drawn with precision Q_ss and then correlated across future times using chol(Q_uu), which gives precision Q_uu ⊗ Q_ss.
  • The factorization cost drops from an (S·T·C)-dimensional matrix to Q_ss (S × S) plus Q_uu (future T·C).

3. 8115946: Use report-only TMB objects

  • tinyVAST() is now called with run_model = FALSE, so it only assembles the inputs. newobj$obj is then built with MakeADFun() without random =. This skips the Laplace setup and inner optimization, which report() does not need.
  • predobj is built with type = "Fun", so only the function value is taped (no derivative tapes).
  • Building the objects is cheaper, and each report() call is a plain function evaluation.

4. deb8287: Pass full unmapped parameter lists to the report objects (mostly correctness)

  • Both report objects are now built without map, so report() takes the full expanded parameter vector, built with unlist(new_parlist).
  • This fixes cases where mapped or shared parameters, or effects fitted as fixed parameters, were not filled in correctly.
  • As a side effect, the per-sample name matching (setdiff / %in% and the loop over tmb_random) is gone. par_template and same_vars are computed once before the loop, which saves a little more time.

As an aside, I believe project() isn't correct for gmrf_parameterization = "projection" yet, but I assume that's quite rarely used at this stage. At most, an error message could be warranted.

@Arnaud-Gruss

Copy link
Copy Markdown

Hi @seananderson and @James-Thorson,

My apologies for testing this only now. For my workflows to work, when I re-install tinyVAST on my laptop, I also need to re-install it on my HPC system, which takes some time. Therefore, I was waiting for Jim to confirm he was happy with the dev-speedups version of tinyVAST before re-installing the dev version of tinyVAST.
I just re-installed the dev version of tinyVAST and re-did probabilistic forecasts for Chatham Rise hake and ... wow ... it took a ridiculously small amount of time to generate these probabilistic forecasts with the conditional approach and 100 parameter draws. I do confirm that the results are identical to the ones I got with the previous dev version of tinyVAST, but the time gained with Sean's developments is just fantastic.
So, thank you so much Sean for these awesome developments and Jim for merging Sean's pull request. Drinks are on me the next time I meet you guys in person :)

Cheers,
Arnaud

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants