[Tracking] Behavioral Foundation Model for Dynamic Foraging — analysis roadmap

[hanhou]

Behavioral Foundation Model for Dynamic Foraging — master tracking issue

Umbrella issue tracking the prioritized analysis roadmap for the disRNN/GRU behavioral foundation model. Child issues below are grouped by priority, mirroring the Prioritized analysis TODO (project plan) in the planning doc.

📄 Live planning doc: https://docs.google.com/document/d/1Xk4Zi9QtQcUNJs4SMvZEb_HbGq0rU8wW3NT8OyNaet4/edit?usp=sharing

Related repos: aind-disrnn-wrapper · aind-disrnn-dispatcher · aind-disrnn-result-access | Project board: org project #184

Target: manuscript by June 2027.

---

📊 Status — 2026-07-12 (4 / 17 original closed · 2 new children added)

Four committed studies now exist under studies/:

study	question	status
01-gru-scaling-law	does more mice → better prediction of unseen mice?	✅ r1–r9
02-gru-scaling-law-ignore	does a headroom-ier (3-way, incl. ignore) target scale?	✅ 48/48
03-disrnn-beta-scan	can we force the disRNN interaction bottleneck to sparsify?	✅ 43/48 + suppl.
04-gru-vs-disrnn-embedding-recovery	do embeddings recover known generative structure?	✅ r1–r2

The one-paragraph scientific state. On 2-way L/R choice the model is near a predictability ceiling: held-out-mouse likelihood saturates by ~100 mice, and the population mean already predicts a new mouse to within 0.3% of full adaptation. But the foundation-model claim survives, on three legs: (1) the population GRU beats a per-mouse classical RL baseline by +0.0136 at D=614 on 100% of held-out mice (p3e-26) — the dominant signal; (2) on the 3-way ignore target likelihood keeps climbing with capacity and never plateaus (H=256/D=614 = 0.7315), so the ceiling was a property of the metric, not the model; (3) on synthetic ground truth, embeddings recover true generative parameters and model family at 97.5–100%, exactly where a correctly-specified baseline breaks down. The interpretable disRNN replicates that recovery for a ~4–6 point likelihood cost — interpretability is nearly free.

The two gaps that most threaten the story:

1. Complete baseline models incl. hierarchical Bayes #20 — no hierarchical-Bayes population baseline. The RL baseline is fit per-mouse independently, so it has no D-axis. The GRU's headline win is partly a "population vs per-mouse" win by construction. This is the fair-comparison baseline and it is not built.
2. Large-scale disRNN/GRU training across full dataset (600–800 mice) #16 — no full-scale disRNN. Every interpretability conclusion (03) is at D=100. The disRNN has never been trained on the full 600–800-mouse dataset.

Where P1 stands: essentially not started. #24/#26/#28/#29 have no code. The chain has a clear entry point — #22's missing logistic-regression readout blocks #26 and half of #29 — and the cheapest first win is #27/#28, where embedding_space_analysis.py already plots subject/session embeddings against metadata and just needs pointing at the D=614 checkpoint.

New modeling direction (#57). A committed design note — docs/design-hierarchical-vi-foundation-model.md, #54, not yet implemented — reframes the FM as an explicit hierarchical mixed-effects model trained by amortized VI (VAE-style). Its key claim: the cognitive hierarchical-Bayes model (#20) and the VI foundation model are the same statistical object, differing only in whether the per-unit latent is hand-specified cognitive parameters or a learned RNN latent. #57 is therefore the other arm of #20, not a replacement for it — and the two must share a held-out matrix and metric from the start, or they can't be compared later without a re-run.

---

P0 — Critical: model training, development, pipeline

• Large-scale disRNN/GRU training across full dataset (600–800 mice) #16 🟡 Large-scale disRNN/GRU training across full dataset (600–800 mice) — GRU done @ D=614; disRNN full-scale missing
• Data split & cross-validation: held-out-subject transfer at scale #17 ✅ Data split & cross-validation: held-out-subject transfer at scale
• Hyperparameter sweep for GRU / DisRNN #18 🟡 Hyperparameter sweep for GRU / DisRNN — N×D + β grids done; GRU HPs never tuned
• Scaling law of performance (model capacity × dataset size) #19 ✅ Scaling law of performance (model capacity × dataset size)
• Complete baseline models incl. hierarchical Bayes #20 🟡 Complete baseline models incl. hierarchical Bayes — Bari baseline done; hierarchical Bayes = top P0 gap
• Synthetic Q-learning population + recovery matrix #21 ✅ Synthetic Q-learning population + recovery matrix
• Generative evaluation suite (run-length, switch prob, logistic coeffs) #22 🟡 Generative evaluation suite — switch-prob + history done; run-length dist. + logistic coeffs missing
• Include-ignore-trials model version (3 actions, engagement state) #23 ✅ Include-ignore-trials model version (3 actions) — engagement state split out → Explicit engagement state (latent engaged/disengaged) — break the ignore-recall ceiling from study 02 #58
• Hierarchical (mixed-effects) foundation model via amortized VI — VAE-style arm of the hierarchical-Bayes comparison #57 🆕 Hierarchical (mixed-effects) FM via amortized VI — design committed, not implemented; the VAE-style arm of Complete baseline models incl. hierarchical Bayes #20
• Explicit engagement state (latent engaged/disengaged) — break the ignore-recall ceiling from study 02 #58 🆕 Explicit engagement state — break study-02's 0.47 ignore-recall ceiling

P1 — Model & embedding analysis (interpretability, mechanism)

• GRU state space & dynamics: fixed-point analysis + embedding reshaping #24 🔴 GRU state space & dynamics: fixed-point analysis + embedding reshaping — not started (unblocked)
• DisRNN update-rule interpretation + embedding modulation #25 🟡 DisRNN update-rule interpretation + embedding modulation — sparsity solved; interpretation not started
• Embedding → policy map #26 🔴 Embedding → policy map — blocked on Generative evaluation suite (run-length, switch prob, logistic coeffs) #22
• Embedding vs. metadata / Q-learning parameters #27 🟡 Embedding vs. metadata / Q-learning parameters — tooling exists, no result yet
• Embedding: within-subject learning trajectory #28 🔴 Embedding: within-subject learning trajectory — not started
• Discover cognitive mechanisms (learning-rate / uncertainty modulation) #29 🔴 Discover cognitive mechanisms (learning-rate / uncertainty modulation) — blocked on Generative evaluation suite (run-length, switch prob, logistic coeffs) #22 + DisRNN update-rule interpretation + embedding modulation #25

P2 — Extensions

• Neural correlation with learned cognitive variables #30 🔴 Neural correlation with learned cognitive variables
• Mixed sim + animal training + deviation analysis #31 🔴 Mixed sim + animal training + deviation analysis — synthetic half already built in 04; cheaper than its P2 label
• Cross-species comparison / transfer #32 🔴 Cross-species comparison / transfer

---

19 child issues (17 original + #57, #58) · priorities mirrored on the project board's Priority field.