A technical introduction for computational neuroscience and bioinformatics researchers.
Brain MoE-PINN is a species-conditioned model family spanning approximately 2.8M parameters for the current C. elegans profile to approximately 1.41B for the current human profile. It learns to generate brain dynamics from multimodal neuroimaging data — EEG, fMRI, and optionally MEG — while using physics-inspired inductive biases and diagnostics rather than claiming a thermodynamically valid model. Unlike a conventional encoder that maps brain data to a static embedding, this model learns a velocity field over a latent state space: given the current brain state, it predicts how the state will evolve. This lets it both represent brain activity (like a foundation model) and simulate its temporal evolution (like a biophysical model). The architecture sits at the intersection of representation learning and dynamical systems theory, targeting a regime — mesoscopic whole-brain dynamics — where neither purely biophysical simulation nor static representation learning is sufficient.
Two research traditions dominate computational brain modeling, and they rarely speak to each other.
Brain simulation solves differential equations for each neuron or neural population. Projects like The Virtual Brain (TVB) use connectome-based neural mass models at 200–1000 brain regions, supporting in-silico lesion experiments and inference via Dynamical Causal Modeling. At the extreme, spiking network simulations on supercomputers (Hygon Exascale, K Computer cerebellar models) integrate Hodgkin-Huxley equations for billions of neurons. These models are generative — you can run them forward without external input — but their parameters are hand-tuned or sampled from priors. They do not learn from data.
Brain representation learning trains encoders that map brain recordings to a latent space. BrainLM (2024) applies masked autoencoding to 6,700 hours of fMRI; NeuroSTORM (2026) scales this to 50,000+ subjects using a Shifted-Window Mamba backbone. For EEG, Magi provides a BERT-style encoder with 3D electrode embeddings. These models support decoding and transfer learning, but they do not generate dynamics — they produce a static embedding for each time window, with no notion of temporal evolution.
Brain MoE-PINN bridges this gap. It combines an encoder stack (representation learning) with a structured stochastic velocity field (simulation), enabling both encoding and generation in a single model. The velocity field uses GENERIC-inspired algebraic candidates—antisymmetry, nonnegative mobility, and pointwise degeneracy projections—as inductive biases. These do not by themselves establish a valid Poisson bracket, stochastic entropy production, or physiological energy/entropy interpretation.
The model processes EEG and fMRI simultaneously, with optional MEG — EEG (19–128 channels at 256 Hz), fMRI (400 brain regions at ~0.5 Hz), and MEG (306 channels at 256 Hz, use_meg=True) — through a pipeline with five stages:
EEG → Magi encoder ──────────┐
fMRI → NeuroSTORM encoder ───┼──→ Hub Token Fusion → z_global
MEG → BERT-medium encoder ───┘ (optional) │
▼
Velocity Field (GENERIC + MoE)
│
▼
Decoder → EEĜ, fMRÎ, MEĜ (if enabled)
EEG — unified Magi backend. eeg_backend="v1" retains the 8-layer,
512-dimensional BERT-medium-style path. eeg_backend="v2" selects the native
24-layer, 1024-dimensional Magi path with RoPE, GeGLU, pre-norm blocks,
alternating global/sliding-window attention, channel-type embeddings,
BIOT-style arbitrary electrode positions, variable channel counts capped by
max_channels, ECoG/sEEG amplitude normalization, and optional Mamba-2
context expansion. Both backends are selected through canonical
BrainMoEPINN/ExperimentConfig; there is no separate model class.
fMRI — NeuroSTORM (frozen backbone). A Shifted-Window Mamba backbone pretrained on 28.65 million frames from 50,000+ subjects. We freeze the backbone and train only a small adapter projection (~5M parameters). Task-specific Prompt Tuning (TPT) prepends learnable prompt vectors to the token sequence, enabling task-conditioned processing without modifying the backbone weights.
MEG — Shared BERT-medium backbone (optional). No pretrained MEG foundation model exists, so we share the EEG encoder's architecture and initialize from EEG weights via load_from_eeg(). This is justified by Maxwell's equations: EEG and MEG both arise from the same post-synaptic currents, differing only in measurement physics (volume conduction vs. magnetic induction). Sharing the backbone enforces this physical correspondence. MEG is disabled by default (use_meg=False); when enabled, the hub fusion expands from two to three tokens and the decoder gains a MEG reconstruction head.
Each encoder produces a sequence of token embeddings. Two learnable hub tokens — one per active modality (EEG and fMRI by default) — attend to these sequences via cross-attention, producing modality-aligned representations. A final fusion step produces a single global latent state
This is the core of the model. Given a latent state
GENERIC-inspired inductive bias.
The learned potentials are evaluated by their observable consequences—reconstruction, forecasting, free-run statistics, state transitions, and external behavioral alignment—not by assigning them physical units or physiological meaning.
The MoE contribution is an optional learned residual, not a universal species default. The C. elegans and zebrafish profiles disable it because their nervous systems are much smaller; mouse and human retain it for the higher-capacity end of the ladder. When enabled, eight shared experts — 2-layer residual MLPs at 1.0× width, ~2.1M parameters each at latent_dim=1024 — provide a dense backbone. Six routed experts (0.7× width, ~1.5M each) are selectively activated via a Top-3 gating mechanism, giving the model capacity to learn specialized dynamical modes. The measured default human MoE velocity block is ~26M parameters. An opt-in deep-expert configuration (--use_deep_experts) replaces the shared stack with 50-layer residual MLPs (52.5M each) and the routed stack with 14-layer MLPs (7.6M each), raising the parameter count sharply.
The router is a recurrent selective state-space model (a Mamba S6 core with 64-dimensional hidden state), not a stateless MLP. This is important: expert routing at time
MoE scope: the structure claim holds for the backbone, not for the default experts. With the default (use_generic_moe=False) experts add an unconstrained velocity after the degeneracy projections, so the total field is not GENERIC even though moe_velocity_share reports the experts' fraction of the applied velocity magnitude and moe_energy_alignment reports how strongly that unstructured part couples to use_generic_moe=True makes each expert perturb delta_L = UJU^\top, delta_M = diag(softplus) + VV^\top, both projected), so the composite field keeps the structure; configs/species/mouse.json already enables it. Historically that mode cost 14.4× the plain experts (379M vs 26M parameters at latent_dim=1024); a bottleneck on the low-rank factor maps brought it to 3.7× (96M), which is why it is now a practical option rather than an ablation curiosity.
The low-rank factorization materialize_poisson=True restores it for diagnostics that need the matrix.
The velocity field operates at multiple timescales through three mechanisms:
Multi-Time-Scale KDA (MT-KDA) — not to be confused with the decoder's Kimi Delta Attention — maintains three parallel exponential moving averages of the velocity with per-step decay rates latent_dt is supplied these are converted to physical time constants
Slow Manifold Projector extracts sub-0.25 Hz ultra-slow dynamics via a learned 256×1024 projection. This is where cross-modal alignment between EEG and fMRI lives — the fast electrical signals and slow hemodynamic signals meet on this manifold through a physics-structured Latent HRF Bridge.
OU-Structured Noise and the SDE policy. Colored (Ornstein-Uhlenbeck) noise is applied according to an explicit noise_mode policy — off | rollout | train | always (default off: deterministic latent steps). The OU state uses the exact discrete recursion integration_dt, derived from latent_dt) by construction. When enabled, the Wiener Homeostat scales calibrate_dim_gain).
Low-rank window SDE (generic path, opt-in). Set features.transition_mode="sde" to replace deterministic Euler updates with Euler–Maruyama over the window clock: diffusion_rank, diffusion_floor, and diffusion_scale configure it. These fields are passed from ExperimentConfig.model_kwargs() to the model. Default transition_mode="ode" remains deterministic and registers no diffusion parameters. Legacy noise_mode controls only OU noise; SDE runs use noise_mode="off" to avoid double-counting noise.
features.scale_anchor="off"|"context" independently controls observation-scale anchoring; it does not depend on ODE/SDE mode. The C. elegans ODE baseline (configs/species/c_elegans_unified.json) and SDE profile (configs/species/c_elegans_unified_sde.json) both select "off". The SDE component pilot uses configs/data/c_elegans_toyoshima_component_sde_pilot.json; its 8-sample generic rollout has forecast_crps=1.0 and a low-weight (forecast_center=0.1) Huber center term. Ensemble-mean component, variance, autocorrelation, correlation-difference, and reconstruction terms are disabled on sampled SDE forecast outputs. Training fails before starting if an active SDE forecast phase has no positive sampled-distribution loss. The matched near-zero control is configs/species/c_elegans_unified_sde_nearzero.json (diffusion_floor=1e-8, diffusion_scale=0). configs/data/c_elegans_toyoshima_component_sde_diffusion_only.json defines a CRPS-only P6 phase; initializing from the ODE P6 checkpoint with --trainable-parameter-prefix transition_diffusion. freezes every parameter outside the diffusion head.
For calibrated forecasts, enable features.use_generic_observation_adapter=true and set positive loss_weights.forecast_crps on forecast phases (or set positive forecast_crps in the data profile; it applies only to selected phases with loss_weights.forecast > 0). The trainer scores {modality}_recon_samples (shape (S,B,K,C,T)) with masked, horizon-weighted empirical CRPS. The pilot uses S=8; finite-ensemble empirical CRPS is a biased estimate of population CRPS. This requires explicit future targets; standard forward() is not on this SDE path. Existing ODE profiles remain unchanged. An older ODE checkpoint can initialize an SDE model; only the diffusion head starts fresh, with empty optimizer slots.
The free-run evaluator reports conditional transition-mean metrics separately from sampled-path metrics. SDE TDE-RICA W1/KL compare real features with pooled projected samples; mean-path W1/KL and D_kernel remain separate. Raw r_sigma, between-sample variance, within-path temporal variance, and raw realization-level lag-1 diagnostics distinguish scale, ensemble spread, and temporal variation. Each horizon reports r_sigma, W1/KL, and D_kernel; SDE transition diagnostics include drift RMS, diffusion RMS, and trace(Q). ODE evaluation uses one path, so mean-path and pooled-sample metrics coincide. Pass --seed (default 42) to reproduce model initialization and sampled paths. The transition-mean path is not the exact multi-step expectation of a nonlinear SDE.
SDE scope limits. The latent endpoint remains window-level; the current C. elegans profile advances 15-second windows. This iteration adds no posterior/SVI uncertainty, semigroup constraint, or finer latent substeps. configs/data/c_elegans_toyoshima_component_sde_path_calibrated.json retains CRPS at lower weight and adds a fixed 32-dimensional random orthogonal channel projection with lags 1, 2, 4, and 8 differences, sampled-path autocorrelation calibration at those lags, and a one-sided pathwise temporal-spread penalty from H5 onward. TDE-RICA remains evaluation-only. S=8 paths are Monte Carlo realizations, not independent training replications. Training logs raw forecast terms as loss_forecast_*; colab_sde_path_calibrated_smoke.ipynb reports raw and weighted contributions.
The path-calibrated data profile resolves the ensemble-mean Huber center, correlation-difference, and variance forecast terms to zero; the Colab smoke uses those same forecast weights. Innovation autocorrelation is reported only for estimable transition lags: an eight-transition rollout supports lags 1, 2, and 4, while lag 8 requires at least nine transitions. The sampled observation path objective separately measures lags 1, 2, 4, and 8 across concatenated forecast frames.
TDE-RICA horizons that cannot yield three projected frames are recorded as skipped with a reason; the rest of the evaluation continues. Per-phase parameter-prefix learning-rate scales are preserved in DeepSpeed optimizer groups and through phase schedules.
The evaluator's r_sigma is the average per-realization raw-observation temporal standard-deviation ratio, not latent Var[z_t] or cross-realization spread. between_sample_variance and within_path_temporal_variance use different conditioning axes and are not additive variance components. sde_diffusion_rms is a covariance-rate scale; sde_noise_drift_ratio compares realized step increments.
Stimulus and intervention control conditioning. Loaders may declare control-role modalities. control_reduction="resample" segments a single current window into per-step means; mean collapses the window, last selects its final frame, and peak preserves sparse pulse amplitudes. Multi-horizon future controls are reduced independently within each forecast window. When one 15-second future window maps to one latent transition, resample is still a 15-second mean and discards within-window pulse timing. The path-calibrated Toyoshima profile records this limitation and does not validate control-conditioned dynamics. Randi optogenetic experiments require control values aligned to finer latent substeps; peak alone does not recover pulse timing.
The decoder uses KDA (Kimi Delta Attention), a linear-complexity attention mechanism from Kimi (arXiv:2510.26692). Unlike standard softmax attention, KDA uses a delta rule — the attention output is computed as an incremental update to a recurrent hidden state, with per-key-dimension gating for selective state-tracking. This makes it ideal for autoregressive generation: each new token updates the state in causal_conv1d component requires CUDA, so a fallback decoder (ModalityDecoderRouter) is available for CPU-only development.
The decoder gate concatenates hub tokens from all active modalities. For a model with EEG + fMRI + MEG enabled, this produces a 3072-dimensional gate input (3 × 1024), enabling cross-modal information to influence reconstruction of each modality.
EEG and fMRI are always reconstructed; MEG is reconstructed only when enabled. The MEG reconstruction loss closes a gradient gap that would otherwise leave the MEG decoder supervised only through the cross-modal hub fusion path.
The model operates in two modes:
Perception mode processes real sensor data through the encoders. After predicting
Imagination mode runs the model forward without external input — the velocity field drives the dynamics autonomously. A Counterfactual Tree Search (depth 3, branch factor 4) explores possible futures; a learned trajectory discriminator scores each candidate path and a soft minimum selects one, after which a free-energy-like scalar (compute_efe: the negative of the learned value function, plus goal distance and an explicit control cost) is reported to the action loss. The tree search does not itself optimize EFE, and that scalar is not the variational Expected Free Energy — there is no observation model, no posterior over hidden causes, and no epistemic term. Rollouts run the base velocity backbone only: MoE modulation, stimulus control, species conditioning, and the multi-time-scale filter are not part of the search. The selected trajectory's endpoint then biases the velocity field, creating a closed action-perception loop.
During training, imagination runs on designated steps (controlled by an imagination_interval curriculum) to manage computational cost. Imagination components are gated behind a use_imagination flag — when disabled, they consume no memory.
Hebbian Associative Memory stores latent states via Oja's rule — a biologically plausible Hebbian update that maintains normalized weights (
A replay buffer accumulates imagined states (capacity 256) and provides temporal pairs for replay loss computation. The buffer is cleared alongside other state when reset_history() is called at phase boundaries.
Magi v2 is a native observation backend, not a second Brain MoE-PINN model.
The implementation is split across the vendored magi/ package and
encoders/eeg_encoder_v2.py:
-
The Magi backbone scales to the ModernBERT-large design target (24 layers × 1024 hidden dimensions), with RoPE, GeGLU, pre-norm blocks, alternating global/sliding-window attention, and tiling initialization from the v1 depth.
-
Channel-type embeddings distinguish scalp EEG, ECoG grid, sEEG depth, and unknown channels. BIOT-style position features support arbitrary electrode layouts, while
max_channelsbounds memory use. -
The foundation wrapper also exposes momentum-encoder, masked-prediction, contrastive-projection, and adversarial-subject-classifier components for standalone Magi pretraining; the canonical BrainMoEPINN forward path uses the projected encoder representation.
-
ECoG/sEEG amplitudes are normalized before encoding (
1/20for ECoG grid and1/10for sEEG depth by default). -
Optional Mamba-2 context expansion and pretrained checkpoint loading are propagated through the canonical model lifecycle methods.
data/ecog_dataset.pyhandles DANDI NWB, EDF, and BIDS inputs, including variable channel counts, MNI coordinates, and paired ECoG-fMRI windows.data/ecog_preprocessing.pyprovides resampling, filtering, rereferencing, z-scoring, bad-channel detection, and coordinate extraction. -
data_loader.EEGDenoiseNetDatasetconsumes the public EEGdenoiseNet epoch files from either the checkout root or itsdata/directory. It accepts the published NumPy/MAT formats, keeps deterministic train/validation/test splits, and synthesizes the benchmark's clean-plus-EOG/EMG views using its RMS/SNR convention. The release has 4,514 clean 2-s epochs, 3,400 ocular artifact epochs, and 5,598 muscular artifact epochs; EOG is 256 Hz and EMG is 512 Hz, with clean EEG resampled for the EMG path. The primary clean epoch is used for Magi masked/causal pretraining and the matched artifact-contaminated epoch is the contrastive view; source arrays must be materialized rather than annex pointer files. -
configs/data/eegdenoisenet.jsonis the ready-to-use Phase -1 profile (--phase -1 --eeg_backend v2 --eeg_channels 1 --data configs/data/eegdenoisenet.json); this EEG-only path does not fabricate fMRI targets. -
Human EEG/MEG/fMRI ingestion is manifest-first and local-only. The
data/ingest_human.py bids-sessioncommand discovers BIDS EEG/MEG files, reads EDF/BDF/FIF/BrainVision/EEGLAB/CTF recordings through lazy MNE adapters, extracts atlas-parcel or bounded voxel time series from BOLD NIfTI, and writes the canonical(C, T)ladder plus channel IDs, masks, per-modality rates, and provenance. Existing fMRIPrep or SPM NIfTI derivatives are valid inputs; neither acquisition nor spatial preprocessing is hidden in the command. Example:python -m brain_moe_pinn.data.ingest_human bids-session \ --bids-root /data/openneuro/ds006040 --subject 01 \ --task rest --modalities eeg fmri --eeg-rate 256 \ --atlas /data/atlas.nii.gz --origin OpenNeuro:ds006040 \ --out /data/data_ladder/human_bids
Heterogeneous records use
data/corpus_pipeline.pysource manifests:format: mnefor EEG/MEG andformat: niftifor fMRI. Relative paths are resolved against the manifest, no dataset is downloaded implicitly, andSpeciesSignalDatasetselects physical-second windows at training time. Different EEG/MEG/fMRI clocks remain separate in<modality>_rate_hz;cross_modal_labelmust be explicitly set to1only for verified synchronization and to10BC8 0for an intentional async pair. Install the optional readers only for the formats used (mne-pythonfor electrical recordings andnibabelfor NIfTI). TorchEEG remains optional for dataset-specific EEG transforms; the canonical adapter does not require it because the same profile also handles MEG and fMRI.
The canonical construction path is:
from brain_moe_pinn import BrainMoEPINN
model = BrainMoEPINN(eeg_backend="v2", use_channel_type_embed=True)configs/species/*.json and configs/ladder.json define the species/data
contract. train.py is the only training entrypoint and accepts
either a JSON profile or explicit model flags. The Magi backend, ECoG data
path, MoE routing, context expansion, and multi-step rollout therefore share
one model/configuration surface.
LossWeights (in training/training_phases.py) separates steering terms —
weights
Where
Reconstruction losses (EEG, fMRI, MEG): Standard MSE between decoder output and input signal. These are the primary training signal — all other losses are auxiliary.
For generic signals, an explicit recon_loss_mix remains active alongside
multi-horizon forecast supervision. Its distributional term discourages
correlation-only or low-variance forecasts from becoming the objective's
cheapest solution. The reconstruction boundary additionally enforces the
raw-unit contract: with normalization: "none", no unit-interval or
z-score conversion is inserted after loading. forward_modalities(..., reconstruct=True) matches each measurable decoded channel's standard
deviation to the valid raw context, ignores masked padding, preserves the
predicted mean and temporal shape, and never reads future targets. The
internal linear head may have a different latent amplitude; the exposed
reconstruction is the raw-scale surface.
Structured-dynamics losses:
-
Degeneracy constraint: pointwise residuals in
$L\nabla S$ and$M\nabla E$ . These are enforced by projection, not by penalty, so the residual is identically zero andgeneric_constraintcarries weight 0 in every phase:$(P_S L P_S)\nabla S = P_S L (P_S \nabla S) = 0$ because$P_S \nabla S = 0$ , and likewise$(P_E \mathrm{diag}(M) P_E)\nabla E = 0$ . The weight only becomes live ifapply_degeneracy_projectionis disabled. It remains a local algebraic statement, not a proof of global GENERIC validity or of the Jacobi identity. - Jacobi diagnostic: requires derivatives of a callable Poisson map and explicit state points. It raises an error when those inputs are absent rather than reporting a false score; it remains disabled for the high-dimensional training path.
- Grassmannian regularization: encourages expert attractor subspaces to be orthogonal, preventing expert collapse.
Auxiliary losses:
- NSP (Neurodynamics Statistics Prediction): Predicts band powers and functional connectivity from latent representations, inspired by the DeeperBrain model. Monitor/dormant at weight 0: the current implementation only matches band powers between reconstruction and target, duplicating the band-power loss; re-enable after implementing the latent-to-statistics head.
-
Cross-modal alignment: Cosine similarity between active hub tokens
(EEG/fMRI by default; MEG included when enabled). Every row carries an
explicit
cross_modal_labelsvalue: 1 for synchronized pairs and 0 for intentionally asynchronous pairs. The HRF bridge uses the same labels and excludes asynchronous rows; missing labels fail closed. -
Velocity smoothness: Temporal TV over emitted rollout
delta_z_sequence; it is dormant unless a phase enables it. - MoE load balancing: Encourages uniform utilization across routed experts, preventing a single expert from dominating.
-
Hebbian regularization: Dormant at weight 0; Oja's rule already keeps
$W^\top W \approx I$ , and its input is detached.data(no gradient). The weight norm is logged as a monitor; spectral control happens through the periodic normalization pass in the training loop. - Latent concentration (tsallis): Dormant at weight 0; SIGReg is the chosen anti-collapse term.
- Action loss: Matches predicted EFE to an explicit executed-action utility target; it is not enabled by default because a self-generated EFE is not a supervised action outcome.
- Replay loss: Matches imagined states to an explicit replay target. It does not use the model's engram output as its own target.
-
Intervention response: Runs matched treated and baseline perturbations and
fits their predicted latent effect to an explicit
intervention_target; missing perturbation/effect contracts raise.
Training has two distinct representations: executable phase configurations and older token-ledger hypotheses. The executable CLI schedule is:
| CLI phase | Loop steps | Configured batch | Nominal batch visits* | Role |
|---|---|---|---|---|
| Phase -1 | 10,000 | 32 | 0.32M | Executable Magi v2 masked/causal/contrastive EEG pretraining |
| Stage 1 (P1–P6) | 235,000 | 16 | 3.76M | Reconstruction, two-horizon forecast where declared, routing, and physics-inspired diagnostics |
| Stage 2 | 20,000 | 16 | 0.32M | Labelled cross-modal latent and HRF bridge |
| Stage 3 | 50,000 | 16 | 0.80M | Hebbian memory and online-assimilation hooks; policy losses are opt-in |
| Total | 315,000 | — | 5.20M | Current executable phase schedule |
* Nominal batch visits are loop steps × configured phase batch; they are
not token counts. DeepSpeed gradient accumulation, sampler reuse, and dataset
windowing determine the effective global-sample count.
--phase_steps N caps every selected phase at N loop steps (a phase already
below the cap keeps its declared budget, and warmup_steps is clamped to the
cap). Bounded-corpus experiments use it instead of signalling a preemption
checkpoint to stop a full-scale schedule; the declared table above remains the
full-scale target.
Stage 2 expands its model context from 4,096 to 8,192 to 16,384. Stage 3
expands from 16,384 to 32,768 to 65,536. These schedules update model-side
context buffers; they do not automatically change the seq_seconds
windows constructed by the data loader. A valid long-context budget therefore
requires an explicit data-window contract and a measured batch trace.
The previous 67B/364B token plans and the 74-day estimate are retained only as historical arithmetic comparisons in §4.4. They are not validated execution targets.
MoE-from-start remains an architectural hypothesis, not a measured scaling result. The current configuration starts with shared and routed experts to avoid a dense-to-MoE conversion, but the earlier 16B/140B token comparison has not been established by convergence or throughput measurements. Treat expert specialization, active-parameter counts, and data requirements as validation questions on held-out corpora.
- Local development: CPU-only development is supported for tests, manifest inspection, and preprocessing. No local GPU throughput is assumed.
- 16 GiB cards: candidate tier for C. elegans, zebrafish, and mouse short-window runs with microbatch 1; peak memory must be measured on the target topology.
- 24 GiB cards: animal-profile pilot tier and human short-context ZeRO-3 smoke tests across multiple cards.
- 32 GiB cards: human short-context calibration tier; do not infer 16k or 65k feasibility without a real batch trace.
- 48–80 GiB cards: candidate tier for human medium/long-context experiments; 65,536-context training remains unclaimed until activation and communication peaks are measured.
- DeepSpeed ZeRO-2/3: use the checked-in configs as starting points, with activation partitioning enabled and CPU checkpointing/offload disabled by default.
- Mixed precision: FP16 autocast is configured with selective FP32-sensitive operations. The configuration keeps FP32 master weights, so parameter-only memory estimates must not be treated as FP16-weight estimates.
- Device-agnostic runtime: all device references are centralized in
runtime/device_utils.py; there is no hardcoded.cuda()path outside that module.
The 67B and 364B token ledgers are historical planning hypotheses. Current parameter counts from the active species configurations are:
| Profile | Parameters | FP16 export | Static training-state lower bound |
|---|---|---|---|
| C. elegans | 2.81M | 0.005 GiB | 0.042 GiB |
| Zebrafish | 14.78M | 0.028 GiB | 0.220 GiB |
| Mouse | 60.33M | 0.112 GiB | 0.899 GiB |
| Human core | ~1.411B | ~2.63 GiB | ~21.03 GiB |
The static lower bound uses approximately 16 bytes per parameter for FP16 model weights, FP32 master weights, gradients, and Adam states. It excludes activations, second-order physics graphs, temporary tensors, communication buckets, allocator fragmentation, and checkpoint duplication. Because the checked-in configuration keeps FP32 master weights and may retain FP32-sensitive live tensors, the actual static term can be closer to 20 bytes per parameter before those overheads.
For an order-of-magnitude parameter-only comparison:
where
| Position ledger | Parameter-only lower bound | Time at 1 effective PFLOP/s | Status |
|---|---|---|---|
| Historical 67B ledger | 0.567 ZFLOP | 6.57 days | Arithmetic comparison only |
| Historical 364B ledger | 3.082 ZFLOP | 35.67 days | Arithmetic comparison only |
| Historical 74-day claim | 7.417 ZFLOP | 85.84 days | Not reproduced; not a budget |
The old 64×V100 assumption (125 TFLOP/s × 14.5% effective MFU) corresponds to approximately 1.16 effective PFLOP/s. At that rate, the 364B lower bound is approximately 30.75 days, so the 74-day claim contains an unmeasured approximately 2.41× overhead factor. Replace all such estimates with measured FLOP/update × updates after a representative cluster calibration.
The first intervention ladder uses processed traces and explicit event metadata rather than raw imaging movies:
| Corpus | Current measured source size | Intervention scope | Status |
|---|---|---|---|
| Local C. elegans salt pilot | 0.197 GiB | Salt stimulus; not optogenetic | Ingested and verified |
| Randi C. elegans PumpProbe, OSF e2syt | 1.2 GiB local export, 113 recordings | Recording-local neuron-targeted optogenetic stimulation; source target IDs, labels, and event files are preserved | Canonical ladder ingested and model-smoke verified |
HuggingFace celegans_neural_data |
669.5 MB parquet; 919 emitted worms | Calcium-only activity; source has no unified intervention track | Downloaded and ingested into the canonical ladder |
| Unified C. elegans manifest root | Small manifest; references the three canonical roots | One stimulus control contract, source-local calcium identities |
1,056 rows; P3-P6 remote run complete with per-source free-run evaluation (§C. elegans unified P3-P6 run) |
| DANDI:001569 | 2.833 GiB, 13 NWB assets | Targeted two-photon photostimulation of rsChRmine-expressing neurons | Inspect one asset before full download |
| OpenNeuro ds001541 | 6.947 GiB, 597 snapshot files | DRN population optogenetic fMRI | Use population/region scope, not neuron one-hot |
| Velez-Angel et al. zebrafish lateral-line study | No public archive identified | Single-neuromast stimulation plus whole-brain calcium imaging | Contact-gated; not included in the storage budget |
The checked-in Randi adapter consumes the local text export without a hidden download:
python -m brain_moe_pinn.data.ingest_randi \
--input ../exported_data \
--output ../data_ladder/c_elegans_randiUse configs/species/c_elegans_randi.json with
configs/data/c_elegans_randi.json for the 2 Hz, 15-second-window,
recording-local target-position experiment. The adapter follows the supplied
TSMixer-Ext pulse levels, uses the measured 0.5-second clock rather than the
gist's 0.05-second plotting value, masks invalid fluorescence frames instead
of deleting target channels, and preserves the original labels under
randi_labels/. The default [0, 200] + per-recording min-max path is the
legacy model-input representation, not a physical fluorescence calibration.
Use --scale prefix_minmax --fit-frames N to fit the model scale on a
calibration prefix without using future frames, or --scale none --max-value ... for source-value analysis. Neither path recovers absolute
voltage/current without external indicator and membrane calibration. It does
not import gKDR-GMM metadata; that metadata remains owned by the separate
Toyoshima salt adapter.
The unified C. elegans profile is
configs/species/c_elegans_unified.json plus
configs/data/c_elegans_unified.json. It keeps source-qualified subjects and
recording-local channel ids, uses batch size 1 with 15-second physical windows,
and presents stimulus at perturbation width 282: feature 0 is the salt or
opto drive, features 1–281 are Randi recording-local target gates, and
HuggingFace rows omit control. The manifest-federated root avoids another
multi-hundred-MB copy; package the three source roots together or run the
merge command with --materialize for a self-contained remote bundle.
The four locally measurable sources above total approximately 10.835 GiB raw. A conservative raw + conversion + canonical cache factor of 3 gives approximately 32.5 GiB for a sequential working set; use at least 100 GiB of cluster scratch for the three-species pilot, including logs and checkpoints. A 25 GiB local volume is not sufficient to stage all sources safely. Collections above 100 GiB should be streamed or staged on at least 500 GiB of scratch/object storage.
For the current human parameter count, static ZeRO estimates are:
| Data-parallel world | ZeRO-2 static estimate/GPU | ZeRO-3 static estimate/GPU |
|---|---|---|
| 4 cards | 9.20 GiB | 5.26 GiB |
| 8 cards | 7.23 GiB | 2.63 GiB |
| 16 cards | 6.24 GiB | 1.31 GiB |
These values omit activation and communication peaks. A human run is considered cluster-ready only after a 1,000-update calibration records peak allocated/reserved VRAM, updates/s, measured FLOP/update, communication wait, and held-out-subject intervention metrics. The resulting budget is:
No VRAM tier currently carries an unconditional 65,536-context guarantee.
The unified profile was measured with
configs/species/c_elegans_unified.json (latent dimension 192,
perturbation width 282, generic observation adapter, no MoE):
| Measurement | Result |
|---|---|
| Total model parameters | 2,860,753 |
| Dynamics parameters | 1,661,258 |
| FP32 parameter storage | 10.913 MiB |
| Largest tested calcium window | 281 channels × 30 frames (Randi, 15 s) |
| Peak incremental CPU RSS after model construction | 39.469 MiB |
The CPU benchmark completed finite forward/backward passes for the largest Toyoshima, Randi, and HuggingFace channel-count rows. The current host inventory reports an 8 GiB unified CPU/GPU memory limit and one Apple A18 Pro integrated Metal candidate; PyTorch reports MPS available, but no dedicated VRAM amount exists and the full benchmark below was run on CPU. The parameter/activation measurement supports a 4 GiB dedicated-GPU planning floor for this exact batch-1 profile, with 8 GiB recommended for allocator, checkpoint, and runtime headroom; 12–16 GiB is the comfortable tier if windows, workers, or future model features grow. These are planning tiers, not a claim about a remote provider's SKU. A remote run still needs a short calibration for peak allocated/reserved VRAM and throughput before scaling.
Remote training has since run at this profile: the unified C. elegans P3-P6 run executed 10,000 loop steps at batch 1 on one RTX 3080 Ti with a peak GPU allocation of 0.4 GiB and 6.6 steps/s, which is consistent with the planning floor above and leaves the 8 GiB tier as headroom rather than a requirement (§C. elegans unified P3-P6 run).
For biological free-run evaluation, use the neighboring ../TDE-RICA toolbox
as an independent reference rather than treating reconstruction loss as a
generator score. The local adapter is tools/tderica_free_run.py; it reports
the native correlation/variance/autocorrelation suite plus TDE-RICA time
alignment, distribution, and optional D3 continuous-dynamics metrics. The
physical clock is mandatory:
python tools/tderica_free_run.py \
--real real_window.npy \
--generated generated_window.npy \
--dt-s 0.333 \
--tderica ../TDE-RICA \
--no-d3 \
--output tderica_report.jsonFor the remote GPU run, first generate deterministic, unforced model output
from the best checkpoint, save (T,N) real/generated windows, and run the
same evaluator in a clean process. Use --full-d3 only after the fast report
is finite; TDE-RICA's transfer-entropy and local-Jacobian estimators are
high-dimensional diagnostics, not training objectives.
manifest-federated root must be transferred with all three canonical source
roots, or materialized with data/merge_c_elegans.py --materialize.
The cross-modal inventory is organized around:
- ECoG/iEEG: AJILE12 (
DANDI:000055), synchronized iEEG-fMRI (DANDI:000623), scalp+iEEG (DANDI:000574), and PtNRGrids (DANDI:000465/000554). - fMRI: UK Biobank, HCP, ABCD, and HBCD, with HBCD treated as separate protocols rather than simultaneous EEG-fMRI recordings.
- Low-resolution EEG: TUH EEG, LEMON, and MPI-Leipzig.
Dataset sizes and token/position counts remain acquisition estimates until the corresponding files are downloaded, standardized, and measured locally or on cluster scratch.
Training this model family — from 2.8M/14.8M/60.3M animal profiles to the approximately 1.41B human profile — with a multi-term steering loss stack, monitored auxiliary terms, second-order physics graphs, and physical constraints requires layered defenses:
Seven-layer stability defense:
- KDA state L1 normalization every 64 steps
- Pointwise degeneracy projection every step
- Kahan summation for floating-point accumulation
- Stable softmax (subtract max before exp)
- Hebbian spectral normalization every 100 steps
- EMA monitoring of observable control proxies and router entropy
- Energy audit: track total energy-like field output for divergence
AutoRollback (training/stability.py) uses explicitly named diagnostics: (1) negative signed-work ratio when a trajectory diagnostic is supplied; (2) energy-like divergence; and (3) router entropy collapse. The first trigger is not an entropy-production estimate. A fourth trigger keyed on a ks_entropy metric was removed: nothing ever emitted that key, so it could not fire.
Wiener monitors run at a fixed cadence (default: every 10 training steps, monitor_interval on BrainMoETrainer) and only in training mode, so evaluation stays history-independent. They regulate stochastic noise only when explicitly enabled; their setpoints are control targets, not universal criticality observables.
During development, several stability-critical issues were identified and resolved:
State management. Two components stored mutable running state as plain Python attributes or nn.Parameter tensors — patterns that fail under .to(device), state_dict(), or distributed training. The PoissonSSMRouter's recurrent SSM state (which accumulates temporal context for expert routing) and the MultiTimeScaleKDA's three timescale states (which maintain fast/medium/slow velocity averages) were both converted to registered PyTorch buffers with in-place .data updates. This ensures they survive device transfers, appear in checkpoints, and are properly synchronized across distributed workers.
Imagination timing. The joint perception-imagination forward pass (which closes the action-perception loop by running CFTS and EFE computation inside perception mode) was accidentally duplicated — the code block appeared twice, causing CFTS to run twice per step. The duplicate was removed, leaving a single coherent path. Additionally, the _imagination_active flag that controls whether imagination runs on a given step was being set after the forward call, creating a one-step lag and incorrectly activating imagination on step 0. It now executes before the forward pass, ensuring the flag reflects the current step.
Dimension hygiene. The action tensor passed to the active inference loop was hardcoded at 2048 dimensions while the latent space is 1024-dimensional. This was fixed to match. In the Poisson router, the standard Mamba S6 discretization
The model makes specific, falsifiable predictions that can be tested through ablation experiments:
| Component | Prediction | Test |
|---|---|---|
| Degeneracy projection | Disabling |
Delta-z Frobenius norm growth > 2× |
| Grassmannian regularization | Without it, expert subspaces collapse (become non-orthogonal) | Off-diagonal Gram norm increases |
| Poisson conservative term | Without it, EEG PSD flattens (loses oscillations) | 1/f slope drops below 0.5 (healthy: 0.8–1.2) |
| SSM router vs. MLP router | Recurrent routing reduces expert switching frequency | Switching frequency < 0.5× MLP baseline |
| Band-power + spectral slope loss | Models without these fail to reproduce correct spectral profiles | Per-band MSE; alpha suppression ratio |
Baseline comparisons against four simplified architectures (VAE+Neural ODE, Graph Neural ODE, No-hub-token fusion, Transformer trajectory model) on five metrics (reconstruction MSE, PSD slope fit, FC correlation, parameter count, throughput) determine whether each complex component provides net benefit.
Several design choices distinguish this project from conventional deep learning architectures:
Forward-increment parameterization. The model predicts a latent increment
Physics-inspired structure, not a theorem. The forward pass combines antisymmetric and dissipative candidate terms, pointwise degeneracy projections, and stochastic-noise modules. Loss terms and diagnostics reinforce these inductive biases, but they do not establish a valid GENERIC or metriplectic system without separate bracket, stochastic-process, and empirical validation.
Emergence paradigm. In Stage 1 P1, all physics losses have zero weight. The model is expected to self-organize through reconstruction pressure and MoE routing competition — complex dynamics should emerge from simple objectives. Physics losses are gradually activated from P2 onward, acting as guardrails rather than hand-holding.
State is explicit, not hidden. All temporal state in the model — the router's SSM state, the KDA timescale states, the Hebbian weight matrix, the replay buffer — is registered as named PyTorch buffers. Nothing is stored as a plain Python attribute that would be lost on .to(device) or checkpoint save. This is a deliberate engineering choice, not just a fix: in a model where temporal continuity is essential, losing state is losing the computation.
What works (verified by smoke tests on CPU; re-verified in the 2026-09 checkout):
- Full forward and backward pass — parameter totals are configuration-dependent. The
magiencoder is vendored in-repo but needs thetransformerspackage (absent from the local dev venv), so end-to-end totals could not be re-measured locally (2026-09); re-measure at first cluster launch. Measured in this checkout: MoE velocity block = 25.6M (default shallow experts) / 466.1M (--use_deep_experts). Older 770M/147M smoke-test totals came from a pre-revision architecture and are not valid resource estimates. - All steering loss terms compute correctly, including the spectral slope loss on reconstructed EEG/MEG; dormant terms are exercised in monitoring/ablation mode (2026-09)
- Gradient flow verified across all components (encoder stack, velocity field, MoE, decoder)
- Router state persists correctly across forward calls and resets
- Multi-Time-Scale KDA states update correctly as buffers
- Spectral slope guard correctly returns zero loss for sequences with fewer than 16 temporal samples
Known limitations:
- The KDA decoder uses
causal_conv1dfrom the Flash Linear Attention (FLA) library, which requires CUDA. A fallback decoder is available for CPU development. - The training loop currently passes
actual_eegas a slice of the input EEG, not data from time$t+1$ . True online active inference requires a paired-sequence dataloader. - Pretrained NeuroSTORM and BrainLM weights need to be downloaded (loading infrastructure is in place).
- Velocity smoothness sits at weight 0 until sequence-level training (multi-timestep batches) emits
delta_z_sequence; it now takes true temporal TV and returns zero for single-step input. The misnamedks/ velocity-difference regularizer was deleted (it was 0.0 in every phase and duplicated this term). - The 1/f constraint now lives in exactly one place (SpectralSlopeLoss); BandPowerLoss's built-in estimator is gated behind
include_one_over_f=Truefor standalone ablation only.
VelocityBrain and the canonical BrainMoEPINN model accept an optional
perturbation/control tensor when constructed with perturbation_dim. This
adds an explicit learned control term to the latent velocity. The interface
supports intervention-response experiments, but conditioning on a perturbation
does not itself identify causality; valid intervention timing, controls,
randomization or another identification strategy, and subject-held-out
evaluation are still required.
diagnostics/causal_dynamics.py provides controlled rollouts,
treated-versus-baseline response metrics, and forward/reverse prediction gaps.
These are operational temporal-asymmetry diagnostics, not thermodynamic
irreversibility proofs.
physics.stochastic_process defines the controlled-SDE contract, Euler-Maruyama
steps, Gaussian transition likelihoods, and conditional forward/reverse path
log ratios. A total path entropy-production value is returned only when the
caller supplies both boundary log densities and a correctly specified reverse
process, including parity handling for odd variables.
Jarzynski/Hatano-Sasa protocol estimators and stronger global Poisson/GENERIC construction are intentionally deferred until controlled datasets and physically anchored state/process definitions are available.
There is one canonical implementation: BrainMoEPINN in __init__.py.
Magi v2, ECoG/sEEG channel metadata, and multi-step rollout are native
capabilities of that model; there is no separate v2 model or compatibility
entrypoint.
Species profiles live in configs/species/*.json, with the ordered replay
contract in configs/ladder.json. Each profile specifies modalities, sampling
units, channel/region bounds, feature gates, and leave-subject-out training
metadata. Species conditioning is metadata conditioning only; it is not a
claim that channels or regions are anatomically homologous across species.
For calcium, voltage, widefield, behavior, and other variable-channel inputs,
call BrainMoEPINN.forward_modalities with a mapping of modality names to
(B, channels, time) tensors. This path runs the shared latent dynamics with
explicit multi-step rollout. It does not fabricate EEG/fMRI reconstruction
targets for species that lack them; instead, with reconstruct=True, any
modality within recon_max_channels is encoded channel-wise
(ChannelSignalAdapter) and decoded back to its raw sampling grid
(SignalReconstructionHead, which generates the waveform from the evolved latent), emitting
{modality}_recon keys. Supervise those with LossWeights.recon_extra
(e.g. {"calcium": 1.0}), which TotalLoss applies through the same generic
reconstruction registry as recon_eeg/recon_fmri/recon_meg. Human
EEG/fMRI/MEG reconstruction remains available through the standard
forward(eeg, fmri, ...) path. A generic profile must not call that path;
generic_observation_only models raise a clear error and require
forward_modalities.
The canonical model lifecycle surface is intentionally small:
- Used by production code:
set_training_step()(encoder freeze/thaw),set_context_length()(progressive context expansion),load_pretrained()(CLI checkpoint loading), andreset_router_state()(phase-boundary state reset). - Retained as useful model operations:
forward_modalities(),reset_history()(clears replay/KDA/active-inference state), andget_num_params()(all-module configuration diagnostics). - Species-capacity reference: use
get_dynamics_num_params()plus the configuredlatent_dim. It counts onlyvelocity_brainand the optionalmoe_velocityresidual that contribute todelta_z; a disabled MoE reports zero parameters. The profiles use latent dimensions 192/384/512/1024 and Poisson ranks 32/64/64/128 for C. elegans/zebrafish/mouse/human. Encoders, observation adapters, fusion, decoders, and active-inference modules are observation transforms or auxiliary consumers, not the species-dynamics capacity. This keeps the dynamics stack dominant in the small-animal and mouse full-model budgets; the full human end-to-end checkpoint also contains a deliberately large observation backbone and decoder, so its raw total is not a valid dynamics-capacity comparison. - Removed after consolidation:
compute_loss()duplicated the trainer'sTotalLosscall and had no caller;BrainMoEPINNConfig.from_dict()and.to_dict()duplicated serialization already owned byExperimentConfig. Loss computation and config serialization now have one authoritative path.
The canonical training entrypoint is train.py; it accepts
--config configs/species/human.json plus explicit Magi v2/ECoG options.
This is a breaking consolidation: old model-module imports and the former
split training entrypoint are intentionally removed.
Hardware planning is profile-dependent:
- CPU-only local development for tests, manifest inspection, and preprocessing; no local GPU throughput claim
- 16 GiB cards for animal-profile short-window pilots
- 24–32 GiB cards for multi-card animal runs and human short-context calibration
- 48–80 GiB cards for measured human long-context candidates
- DeepSpeed ZeRO-2/3 with torchrun elastic training
- No fixed 64×V100 requirement or 74-day full-scale estimate is currently validated
Implemented in this checkout:
- Canonical
BrainMoEPINNconstruction with Magi v1/v2 backend selection. - ECoG/sEEG channel metadata, amplitude normalization, variable-channel handling, and ECoG preprocessing/dataset adapters.
- MoE-from-Stage-0 configuration and unified JSON species/data-ladder profiles.
- Multi-step latent rollout, context expansion, pretrained-weight loading
hooks, and a single
train.pyentrypoint.
Operational work still required before claiming cluster-scale readiness:
- Acquire and validate the first intervention corpora: Randi C. elegans PumpProbe, DANDI:001569, and OpenNeuro ds001541; preserve target IDs, population scope, event timing, and source provenance.
- Download one DANDI:001569 NWB asset and inspect target coordinates, photostimulation intervals, ROI traces, and clock units before staging the full draft.
- Re-measure end-to-end memory and throughput for each target VRAM tier; record peak VRAM, FLOP/update, updates/s, communication wait, and checkpoint size.
- Rebuild the token/position ledger from actual dataloader windows and phase configs; do not use the historical 67B/364B/74-day claims as resource commitments.
- Run multi-node forward/backward, checkpoint-resume, and elastic-recovery tests on the target topology.
- Run Magi Phase -1 against a real EEG loader and verify checkpoint/resume, masked leakage, causal NTP, and EMA updates on the target topology.
Validation is considered meaningful only when it covers reconstruction loss, gradient flow, router behavior, free-run stability, spectral fidelity, cross-modal alignment, parameter/memory budgets, and throughput. The historical compute numbers above are planning comparisons, not results already established by this repository. The architecture summary's predicted gains are hypotheses for those ablations, not results already established by this repository.
This section consolidates the project's literature review and research brainstorm. It is an engineering research note, not a publication-ready systematic review; external claims and quantitative comparisons still require source verification before publication.
Brain simulation builds a generative forward model whose dynamics produce observable neural or behavioral data. It supports inference when observations are used to estimate hidden states and supports in-silico experiments. Brain representation learning maps recordings to latent embeddings. It is useful for decoding and transfer learning, but the encoder alone does not generate neural dynamics. Brain MoE-PINN combines representation-learning encoders with a generative latent dynamical core and active inference.
Spiking simulations solve neuron-level ODEs and communicate through spike events. Their cost scales with neurons and synapses, making them memory-bound and dependent on large supercomputers.
| Project | Scale | Hardware | Model type | Inference | Validation |
|---|---|---|---|---|---|
| BlueGene cortical column (2005–07) | ~1 mm³ cat cortex | IBM BlueGene/L, 128K CPUs | Hodgkin–Huxley spiking | No | Resource metrics |
| K Computer cerebellum (2012–15) | Cat-to-human scale | Fujitsu K, 82,944 CPUs | Izhikevich spiking + OKR | No | Firing-rate comparison |
| Hygon exascale brain (2022–26) | 86B neurons, 47.8T synapses | Hygon prototype, 14,012 GPUs | Conductance-based + Bayesian inference | Yes, EnKF/PF | BOLD comparison |
| PEZY-SC cerebellum (2017) | Human scale | PEZY Shoubu | Izhikevich spiking | No | OKR behavior |
The Hygon approach is particularly relevant because it adds a Bayesian inference layer over a spiking simulation. An observation operator maps hidden neural states to BOLD observations, and a particle filter or Ensemble Kalman Filter estimates the hidden trajectory. This is a useful reference for future latent-state assimilation in Brain MoE-PINN.
Limitations of this class include phenomenological detail at extreme compute cost, hand-tuned or prior-sampled parameters, inference overhead, and limited validation against realistic behavioral tasks.
| Project | Scale | Inference | Validation |
|---|---|---|---|
| The Virtual Brain (TVB) | Whole brain, 200–1,000 regions | DCM and source imaging | Functional connectivity, resting-state networks |
| Dynamic Causal Modelling (DCM) | Small regions, 2–10 | Variational Bayes | Task-evoked responses |
| Jansen–Rit neural mass | Single region to whole brain | Limited | Spectral matching |
TVB is the closest conventional analogue to the project. Both operate at population scale and evolve latent states, but TVB uses fixed biophysical models while Brain MoE-PINN learns GENERIC-inspired potentials from data.
Procedural-connectivity work computes synaptic weights from coordinates and distance rules instead of storing large connectivity matrices. The approach trades memory for computation and is well suited to GPUs. Brain MoE-PINN takes the same philosophy to latent dynamics: its computation is generated from learned energy, entropy, mobility, and Poisson candidates rather than stored as a neuron-level connectome.
| Model | Architecture | Data | Task |
|---|---|---|---|
| BrainLM | Transformer masked autoencoder | Large-scale fMRI | Masked prediction and network identification |
| NeuroSTORM | Shifted-Window Mamba | 28.65M frames, 50K+ subjects | Masked autoencoding, redundancy dropout, prompt tuning |
| fMRI-PTE | Patch-based ViT | HCP | Contrastive learning |
| Magi | 12-layer Transformer, 768d | EEG | Masked prediction, MoCo, PSD, subject invariance |
| BIOT | BERT with 3D electrode embeddings | EEG/MEG | Transfer learning |
| Brain-OF | Flexible-resolution sampling | EEG | Montage-flexible classification |
| DeeperBrain | Transformer + NSP | EEG | Neurodynamics-statistics prediction |
NeuroSTORM demonstrates the value of linear-complexity shifted-window processing, redundancy dropout, and parameter-efficient prompt tuning. Magi is the closest EEG foundation-model reference. DeeperBrain's NSP objective inspired prediction of power spectra, functional connectivity, and cross-frequency coupling from latent representations.
These encoders produce representations; they do not, by themselves, generate future brain activity or support latent-state inference.
| Model | Fusion method | Role in this project |
|---|---|---|
| Brain Harmony | Hub-token cross-attention | Reference for modality-identity tokens |
| CineBrain + CineSync | Multimodal fusion encoder + neural latent decoder | Candidate synchronous EEG–fMRI data and downstream comparison |
Hub tokens carry modality identity while allowing all active modalities to interact in a shared latent space. CineSync reconstructs video from brain signals but does not model autonomous neural dynamics; its temporal structure comes from the video target.
Very few systems bridge representation learning and generative dynamics:
- CineBrain + CineSync provides synchronous EEG–fMRI data and a multimodal decoder, but not a learned brain-state dynamical system.
- Spisak and Friston motivate attractor emergence, orthogonalized representations, asymmetric couplings, and non-equilibrium steady states from Free Energy Principle objectives. This informs the project's Grassmannian attractor regularization.
- Metriplector treats metriplectic dynamics as a computation primitive. Its conservative Poisson and dissipative entropy-gradient split parallels the GENERIC-inspired structure used here, although its scope is broader than neuroscience.
The Brain MoE-PINN architecture is:
flowchart LR
accTitle: Brain MoE-PINN Literature Position
accDescr: EEG and fMRI representations feed a shared latent dynamical system that predicts observations and supports inference and imagination.
eeg["EEG"] --> magi["Magi encoder"]
fmri["fMRI"] --> neurostorm["NeuroSTORM encoder"]
magi --> fusion["Hub token fusion"]
neurostorm --> fusion
fusion --> latent["Shared latent state"]
latent --> dynamics["VelocityBrain<br/>GENERIC-inspired dynamics"]
dynamics --> decoder["Observation decoder"]
decoder --> observations["EEG / fMRI / MEG reconstructions"]
observations --> inference["Active inference"]
inference --> latent
dynamics --> imagination["Free-run imagination"]
The project is a learned generative dynamical system at mesoscopic brain-region scale:
- Encode multimodal observations into a shared latent state.
- Evolve the state with structured stochastic dynamics.
- Correct the state through prediction error and active inference.
- Imagine counterfactual trajectories without observations.
- Learn the observation model and dynamics end to end.
| Dimension | Spiking simulation | Neural mass / TVB | Foundation encoder | Brain MoE-PINN |
|---|---|---|---|---|
| Generative dynamics | Yes | Yes | No | Yes |
| State inference | Some Bayesian approaches | DCM and variational methods | No | Active inference |
| Learns from data | Usually no | Parameter fitting | Yes | Yes, end to end |
| Physical constraints | Biophysical equations | Neural-mass equations | None | GENERIC-inspired algebraic biases |
| Free-run imagination | Yes | Yes | No | Yes |
| Multimodal | Usually limited | Usually limited | Modality-specific | EEG + fMRI + optional MEG |
| Mixture of experts | No | No | No | Yes |
| Online adaptation | Limited | Limited | No | Hebbian memory + active inference |
The closest conceptual analogue is TVB combined with DCM: a dynamical core with inference. The key distinction is that Brain MoE-PINN learns both the dynamics and observation model, scales the dynamics through MoE experts, and adds explicit stability and monitoring mechanisms.
Replace gradient-only active-inference correction with a latent-space EnKF.
Maintain an ensemble of latent states VelocityBrain, decode each member into observations, and update the ensemble
when EEG or fMRI observations arrive.
Potential benefits:
- Per-dimension uncertainty for constrained and imaginative latent directions
- Natural handling of EEG/fMRI temporal-resolution mismatch
- A Bayesian objective based on marginal likelihood rather than MSE alone
The main cost is
Condition the learned Poisson and mobility candidates on a subject-specific
structural-connectivity prior
# Current: operators depend on z only
L_z = poisson_op(z)
# Proposed: operators also depend on structural connectivity
L_z = poisson_op(z, condition_s)A learned adapter could project a DTI tractography matrix into the latent
dimension and inject it into the operator computation. This would support
zero-shot subject personalization, artificial lesioning of
Selective component ablations can test whether model subsystems have interpretable functional roles:
- Drop selected experts and measure brain-state switching
- Remove the antisymmetric
$L(z)$ candidate and inspect dynamics - Ablate salience-gated experts and measure transition detection
- Compare reconstructed spectra, ERPs, band power, and connectivity with lesion-related observations
A concrete prediction is that removing Core Shared experts should disrupt resting-state functional connectivity more than task-specific responses, while removing Specialized experts should produce the opposite pattern.
Replace a dense
where
Use learned conserved-quantity candidates as decoder features:
-
$E(z)$ and$S(z)$ - Casimir-like invariants
-
$|L(z)|$ and$|M(z)|$
For example:
eeg_recon = decoder(z, E=z_energy, S=z_entropy, L_norm=l_norm, M_norm=m_norm)This is a speculative inductive bias, not evidence that the learned fields are physical energy or entropy.
Validate reconstructed EEG with neurophysiological metrics in addition to MSE:
- Delta, theta, alpha, beta, and gamma band-power modulation
- Phase–amplitude coupling
- Spectral coherence
- Task-linked alpha suppression or other stimulus-specific signatures
A possible auxiliary objective is:
Use a two-time-scale assimilation process:
- Assimilate EEG at every fast latent step
- Assimilate fMRI only when a slow observation is available
- Apply hemodynamic-delay handling to fMRI observations
- Free-run when neither observation is available
This would force the latent state to remain stable during fMRI gaps while allowing rapid EEG corrections.
Train variants with and without Grassmannian regularization, then measure whether expert attractors orthogonalize naturally. If orthogonalization persists, the explicit regularizer may be redundant. If it disappears, the result distinguishes an emergent training effect from an explicit inductive bias.
| Idea | Impact | Effort | Novelty | Priority |
|---|---|---|---|---|
| EnKF over latent dynamics | Very high | 2–3 weeks | High | P0 |
| Low-rank Poisson operator | High | 1 week | Medium | P1 |
| Model lesioning | High | 1 week | High | P1 |
| Cross-frequency validation | High | 1 week | Medium | P1 |
| Connectome-conditioned dynamics | Very high | 2–3 weeks | Very high | P2 |
| Multi-rate EEG/fMRI assimilation | High | 2 weeks | High | P2 |
| Stress-energy readout | Medium | 2 weeks | Very high | P3 |
| Free-energy attractor ablation | Medium | 3 days | High | P0 quick check |
The existing decoder implementation is documented in
decoder/kda_decoder.py. The KDA mechanism is a linear-complexity attention
component used for autoregressive reconstruction. XMMM Consciousness Algorithm
is retained as a speculative control-theoretic comparison, not as a validated
model dependency.
Code: development checkout in Google Drive (マイドライブ/a/brain_moe_pinn, self-contained, git-tracked); canonical HPC path /home/yanlu/Documents/a/brain_moe_pinn/.
Fundamental principle: Model velocity $\Delta z$, not absolute position. Irreversibility is a feature, not a bug.
Canonical implementation: one BrainMoEPINN model, one training entrypoint
(train.py), unified JSON experiment profiles. This section is the current
status only (consolidated 2026-09); earlier per-session change logs and the
removed design documents live in git history.
Model and entrypoints
- Single canonical
BrainMoEPINN+ExperimentConfig/BrainMoEPINNConfig;configs/species/*.jsonprofiles,configs/ladder.json,config.pyspecies/modality contracts (metadata conditioning, not anatomical homology claims). train.pyis the only training entrypoint: Magi v1/v2 backends, ECoG datasets, MoE-from-stage-0, Mamba-2 context expansion, species conditioning, pretrained-weight CLI (--eeg_checkpoint,--fmri_checkpoint; files not yet downloaded).- Multi-node launchers at repo root (
brain_moe_launch.sh,brain_moe_multinode.sh,launch_all_nodes.sh), DeepSpeed ZeRO-2/3 configs, phased curriculum and optimizer transitions intraining/training_phases.py.
Encoders
- EEG: unified Magi backend - v1 (8L x 512d) and v2 (24L x 1024d, RoPE,
GeGLU, alternating global/SWA attention, tiling init) with BIOT-style
electrode embeddings, channel-type embeddings (scalp/ECoG/sEEG),
ECoG amplitude normalization, and
max_channelsbounds. - fMRI: NeuroSTORM shifted-window-Mamba backbone (frozen) with TPT prompt
tuning; BrainLM fallback;
load_pretrained()infrastructure. - MEG: shared BERT-medium backbone with EEG,
load_from_eeg()transfer; optional (use_meg). Mamba-2 post-encoder context stacks optional.
Dynamics core
VelocityBrainGENERIC-inspired structured stochastic dynamics: antisymmetric low-rank Poisson candidate, nonnegative diagonal mobility candidate, pointwise degeneracy projections, OU noise with tensor-native Wiener homeostat and per-dimension gains. Antisymmetry is an inductive bias, not a Poisson bracket (see §2.3 for the exact claim boundary).- MoE velocity field: shared + routed experts (deep-expert option), SSM
router with temperature curriculum, single-authority grassmannian
weight on the raw orthogonality emitted by the module, SVD fallback in
orthogonalize(). - Species conditioning, low-rank test-time subject adaptation, slow-gate state transitions, latent diffusion head (non-recursive sampling).
Fusion, decoders, memory
- Hub-token fusion across active modalities, slow-manifold projector + latent HRF bridge (EEG/fMRI alignment on the sub-0.25 Hz manifold).
- KDA decoder router (FLA accelerator with CPU fallback), MEG decoder.
- Hebbian associative memory + replay buffer; counterfactual search and
imagination sampler behind
use_imagination; active inference with Smith predictor and precision gate behinduse_active_inference.
Losses and monitoring (2026-09 architecture)
LossWeightsseparates steering terms (weight > 0) from monitored terms (weight 0; classes kept, metrics logged byTotalLoss._update_monitors). Nothing deleted; see §3 for the full steering set and the rationale per dormant term.- Reconstruction is modality-generic (
ReconstructionLoss.forward) with selectable criteria - mse/l1/huber/poisson/correlation/corr_diff/ wasserstein1 - per modality viaLossWeights.recon_loss_types; species defaults: calcium -> correlation, voltage -> huber. - Spectral constraints operate on reconstructed EEG/MEG only; SpectralSlopeLoss is the single 1/f authority (fMRI excluded, min 16 samples); band power is band-matching only.
- All terms EMA-normalized with frozen calibration (LossNormalizer).
Cross-species signal path and data (2026-09)
forward_modalities()runs arbitrary modalities through the shared dynamics; withreconstruct=Trueit emits{modality}_reconviaChannelSignalAdapter+SignalReconstructionHead(channel-count independent, conditioned on the evolved latent).- Trainer auto-routes generic models through that path and auto-enables
per-modality reconstruction weights and criteria per phase
(
augment_phase_loss_weights,make_dummy_generic_signals). Data plumbing:train.py --data <profile.json>builds real loaders - paired EEG/fMRI/MEG (PairedBrainDataset, MEG optional per stem), manifest-driven species dict batches (SpeciesSignalDataset, contract{modality: (B, C, T)}+ optional masks/metadata), and EEGdenoiseNet clean/artifact view batches for Magi Phase -1. P0-P2 loader layer (2026-09): subject/session/condition manifests, group-preserving train/val/test splits (leave-subject-out), per-channel masks threaded into masked reconstruction, seconds-based windows with time padding, channel-union alignment, load-time region aggregation, trial windows, static graphs, and HDF5/NWB source readers (data/readers.py). MEG batches auto-enablerecon_meg; dataloaders feed train, validation, and held-out test paths. - SDE/control layer (2026-09):
noise_modepolicy (off|rollout|train|always) with a dormant-OUNameErrorfixed (noise was never applied before); stimulus control wired data→trainer→perturbation with zero-init no-op semantics and conditioned gates; roles declared per modality (signal|control|aux|graph);--noise_modeCLI, phase-levelnoise_modeoverride,perturbation_normmetric. - Multi-step rollouts (
num_steps/phaserollout_steps > 1) emitdelta_z_sequence;VelocitySmoothnessLossis true temporal TV over that axis (single-step inputs return zero - legacy batch-axis TV semantics removed). - Time conventions: every time-valued quantity is declared in seconds
(species profile:
sample_rate_hz,sequence_seconds; modules:taus), frame counts are derived (frames = seconds * rate) at the data boundary, and a latent step's physical duration is supplied per batch. One latent step advancesz <- z + dt * delta_z, wheredtis scalar (legacy) or one value per sample; the MT-KDA time constants, the OU transition and the emission filter all evaluate with that duration, so their seconds semantics hold for any sampling rate.latent_dtremains the nominal per-species clock used when a batch supplies none. - Per-sample rates:
species.sample_rate_hz_sourceis"profile"(uniform corpus; the loader validates every manifest row againstsample_rate_hzand raises on mismatch) or"manifest"(rates vary per recording, e.g. per-worm C. elegans imaging; the manifest clock is authoritative). Batches carrydt(seconds per frame) and, when replay mixes species, a per-samplespeciestag that reaches the model's conditioning. - One model per species, one shared code path: a ladder stage runs its own
BrainMoEPINNagainst its own ladder, profile and observation channels. Species are never mixed inside a batch or a model instance -- channel spaces and reporters differ, so a joint batch is not representable; cross-stage continuity comes from checkpoint transfer between stages. - Dynamics -> fMRI (
core/hrf.py,readout: "bold"): the observation channel for BOLD. Friston et al. 2003 (NeuroImage 19:1273-1302, Eq. (3)-(4) + Table 1) supplies the structure and the initialisation -- the Balloon-Windkessel impulse response (kappa/gamma/tau/alpha/rho/V0 = 0.65/0.41/0.98/0.32/0.34/0.02) -- but not fixed physics: our drive is an internal latent with an arbitrary scale, not the DCM's neuronal state, so the emission isreference HRF + zero-initialised smooth correction, fitted with the corpus's own objective. It is causal, TR-aware (the recording's frame interval is the TR; the drive is held across each TR) and reports measured descriptors of the effective HRF (peak latency, FWHM, undershoot ratio, deviation from the reference), which is what makes a human (~5 s peak) and a mouse (~1-2 s) response comparable without trusting the parameterisation.human.json/mouse.jsondeclare it underdata.sensors.fmri; electrical modalities are skipped (their own encoder/decoder carries the physics), and a non-haemodynamic reporter underfmriis rejected. - Zebrafish (
data/zapbench.py): ZAPBench (arXiv:2503.02618) whole-brain light-sheet traces (71,721 cells x 7,879 volumes at 914 ms/volume, nuclear GCaMP7f) mapped to the ladder as one session per stimulus condition (subject= animal,session/condition= gain/dots/flash/taxis/turning/ position/open loop/rotation/dark;taxisis the benchmark holdout), the 26-d stimulus-feature bank as a control-rolestimulusmodality, and--region-bins/--max-cellsreduction because 71,721 cells exceed every channel cap.--rate-hzis required: the volume interval is never assumed. - Ingestion + sanity for real calcium recordings:
data/ingest_c_elegans.py(time-major CSV ->(C, T)npy ladder + manifest) andtools/real_data_sanity.py(real-data forward/backward with species reconstruction criteria, reported next to persistence/channel-mean baselines). The ingest takes per-sample frame rates and the salt-stimulus timing fromstimulation_timing.xlsx(gKDR-GMM), applies the reference channel quality filter (autocorr(lag=20) > 0.3, ~46% of named channels on the pilot), writes the salt drive as a control-rolestimulusmodality plusstimulus_trials/*.json, and records animal/anaesthesia provenance so a re-imaged animal cannot be split across train and val. Verified on the 24-worm salt pilot: 1,625 canonical channels, rates 3.69-5.72 fps, 60 s windows (244 frames at the median rate). - Observation channel (
data.sensors, resolved bycore/sensors.py): every optical modality declares its imaging method (LSFM/SPIM, LFM/XLFM, remote-scanning LSFM, SCAPE/3D-AOD two-photon, spinning-disk 4D, widefield 2P) and its functional reporter (GCaMP6f/6s/7f, jGCaMP8f/s, H2B-GCaMP6s/7f, YC2.60 FRET, Positron2-Kv, Voltron, Arch, pERK). It also declares the model-facing observed space (dff,ratio,standardized,raw_fluorescence,fraction, orvoltage) and whether acquisition/ preprocessing already filtered the target. The latent decoder still applies one model-side latent-to-observed temporal pass; the provenance flag does not add a second filter or disable that pass. - That sensor declaration fixes the frame-integration window, the slice/plane
phase smear, the readout kind and whether windowed dynamics are meaningful
at all (a pERK fixed-tissue map is refused for next-step targets). With
features.use_sensor_emission(on forc_elegansandzebrafish) the decoder emits through that channel -- per-channel indicator low-pass in seconds plus a representation-appropriate readout. Hill saturation is used only for calibrated raw/fraction fluorescence; standardized YC2.60 ratio and dF/F remain in their observed numeric space. GEVI voltage skips the Hill term. The fittedtau/h/Kdare reported with their provenance (imaging,reporter,calibration_required). The latent dynamics are untouched; reconstruction metrics carry their units and their persistence/channel-mean baselines. Control tracks are validated too:control_diagnosticsreportscontrol_active_fracandcontrol_collapse_ratio, and training warns once when a single-step rollout reduces an alternating drive to ~0. Diagnostics, stability, preprocessing - Offline diagnostics:
diagnostics/free_run_metrics.py(multi-axis generative quality),diagnostics/scaling_probe.py,diagnostics/causal_dynamics.py(controlled rollouts and forward/reverse gaps - operational asymmetries, not irreversibility proofs). Trajectory diagnostics report velocity/work/divergence proxies with explicit names. - Seven-layer stability defense incl. periodic Hebbian spectral normalization, L1 KDA-state normalization, router orthogonalization, monitor cadence and eval-determinism (monitors train-mode only), auto-rollback on signed-work criteria only.
- Preprocessing: EEG/fMRI/MEG/paired foundation-mode pipelines plus
MNE-based ECoG pipeline (
data/ecog_preprocessing.py) and ECoG dataset loaders (DANDI/BIDS/EDF).
Verification in this checkout (2026-09)
154 passed, 2 skipped(156 collected). Coverstest_unified_model,test_new_capabilities,test_free_run_metrics,test_signal_reconstruction,test_data_and_rollout,test_manifest_dataset,test_control_sde,test_audit_invariants,test_grad_scaler,test_xla_device,test_deepspeed_config. The same suite passes unchanged on a Kaggle T4, a Colab TPU (v5e1), a MooreThreads MTT S4000 and a Hygon K500SM_AI DCU.tests/conftest.pybootstraps imports; modulespy_compileclean;train.py --helpverified.- On a Kaggle T4 the same suite is
134 passed, 1 skipped(the extra pass is the CUDA-gated GradScaler test).tests/test_smoke.pyneeds the FLA CUDA stack and skips cleanly when absent (pytest.importorskip). - MoE velocity block measured: 25.6M (default experts) / 466.1M (deep). End-to-end parameter/memory totals still require a cluster launch (see open items).
Mixed precision and device portability
- DeepSpeed runs use
torch_autocast, never a nativefp16/bf16block. Native blocks cast module weights to half and disable autocast insideengine.forward(), which this model cannot use: its physics differentiates withcreate_graph=True, returns fp32 gradients, and holds fp32 structural constants.apply_precision()enforces this and strips any native block rather than leaving a contradictory pair. --precision {fp16,bf16}selects the autocast dtype (fp16 default). fp16 and bf16 differ in exactly one field, so they share one config.rectify_deepspeed_config()rebuilds the three batch fields from the real loader batch and world size; before this, no shipped config satisfied DeepSpeed's assertion for any launch topology.grad_scaler()starts atDEFAULT_INIT_SCALE = 4096, not PyTorch's2**16. The opening steps produce the run's largest gradients, so2**16overflowed fp16 on step 1 and the scaler spent a skipped optimizer step halving down. It is returned disabled on backends where loss scaling has nothing to protect.- XLA/TPU: detection is hint-gated (
PJRT_DEVICE,XRT_TPU_CONFIG,TPU_*) because importingtorch_xlainitialises PJRT as a side effect. Autocast on XLA resolves to bfloat16 - TPU has no fp16 compute path, and an explicit fp16 request is promoted with a warning. Verified on a Colab TPU v5e1 (torch_xla 2.9.0): auto-detect resolves toxla, and model forward/backward/step runs clean. - MUSA (MooreThreads MTT S4000, torch 2.2.0 + torch_musa 1.3.0), verified
on hardware: auto-detect resolves to
musa;torch.amp.autocastrejects the backend ("has not registered a module"), so autocast routes totorch_musa.core.amp.autocast;grad_scalerreturns the vendorGradScaler; forward/backward/step runs clean. Vendor scalers are dispatched from_VENDOR_BACKENDS, which covers NPU/MUSA/MLU. - MUSA ships
upsample_linear1dforward only - no backward kernel - soF.interpolatekilled training there._resample_timetherefore uses a hand-rolled_linear_resampleonly whenx.device.type == "musa"; every other backend keeps the tuned native kernel. The fallback is equivalent to 2.4e-07 forward and bit-identical backward. - Hygon DCU (K500SM_AI, ROCm/HIP 6.3 via DTK 26.04, torch 2.9.0), verified on
hardware: HIP presents through torch's cuda namespace, so the existing
cudaentry in_BACKEND_PRIORITYhandles it with no new backend - nohipbranch is needed. Device auto-detect resolves tocuda,grad_scalerreturns an enabled scaler atDEFAULT_INIT_SCALE, and model forward/backward/step runs clean in fp32, bf16 and fp16. DeepSpeedtorch_autocastwas verified end-to-end on this native-bf16 part with both--precision bf16and--precision fp16(ZeRO-2, 6 steps, finite and decreasing) - the combination that cannot be exercised on Turing (no native bf16) or on an unavailable Colab TPU. The DTK environment must be sourced first (source /opt/dtk/env.sh), or torch fails to loadlibgalaxyhip.so. _disabled_scalerfalls back totorch.cuda.amp.GradScalerbecausetorch.amp.GradScaleronly exists from torch 2.3; depending on the newer attribute alone raised AttributeError on the MUSA host's torch 2.2.0, which took out every non-CUDA/XPU device at once.
| Priority | Item | Action |
|---|---|---|
| HIGH | Pretrained checkpoints | Download NeuroSTORM/BrainLM/Magi weights; CLI and loaders are ready |
| HIGH | Dynamics -> fMRI (BOLD) | Implemented as a learnable emission (core/hrf.py): Friston 2003 Balloon-Windkessel response as the initialisation, zero-init smooth correction + learnable vasodilatory lag, gains and compressive output; causal, TR-aware, and reporting the measured effective-HRF peak/FWHM/undershoot plus the deviation from the reference. Wired for human/mouse (data.sensors.fmri). Next: fit it on a real BOLD corpus (Kondo mouse widefield fMRI / HCP-style human runs) and compare the fitted peak latency against the measured HRF |
| HIGH | Targeted intervention corpora | Randi PumpProbe text adapter is implemented and verified on 113 local recordings with 2 Hz timing, target-gated controls, and pulse-preserving reduction; targets remain recording-local positions because source labels are incomplete/duplicated. The Randi intervention evaluation now runs end-to-end (112 rows, drive-aligned windows), and it measures what the model does with the drive, not just what the corpus contains. DANDI:001569 and OpenNeuro ds001541 still need scope-specific adapters and held-out-target validation |
| HIGH | Control pathway supervision | The unified P3-P6 checkpoint is materially unsteered: on drive-aligned windows the recorded drive moves the free run by 0.4 % (salt) to 1.8 % (Randi) of the real effect, with cosine ≈ 0 against both the real effect direction and the 1.68× targeted-channel contrast the Randi data itself carries. perturbation_map starts at exactly zero, so the path only grows from forecast gradients. That checkpoint also predates the future-control alignment fix in partition_generic_batch, so the next gate is a post-fix rerun of tools/intervention_eval.py; only then choose a dedicated intervention loss or drive-weighted sampling. |
| HIGH | Calibrated distributional objective | Held-out W1 0.837 and mean KL 14.72 against gates of 0.05 / 5, with kernel transition 5.26 on the D3 subset; P4-P6 improve every axis over P3 but do not approach the gates. The next experiment is a calibrated distributional objective in the fixed TDE-RICA component space, not more steps of the reconstruction mixture |
| HIGH | Zebrafish (ZAPBench) | Adapter implemented and verified offline against the release constants (data/zapbench.py: condition sessions, 26-d stimulus control, region/cell reduction, mandatory --rate-hz, provenance columns; configs/species/zebrafish.json + configs/data/zebrafish_zapbench.json). Next: point --source at gs://zapbench-release/volumes/20240930/traces (zarr + gcsfs) and ingest a region-reduced subset |
| HIGH | Real corpora ingestion | Salt pilot, Randi PumpProbe, and HuggingFace C. elegans ladders are ingested and verified: salt uses per-sample timing and animal provenance; Randi uses 113 recordings, measured 0.5 s timing, masked fluorescence QC, preserved source labels/events, and recording-local optogenetic controls; HF emits 919 worms from 42,798 rows with unlabeled slots dropped. The unified manifest root and profiles are staged; its control path pads source-width controls to 282. The first unified P3-P6 training run, its per-source free-run evaluation, and the controlled intervention evaluation are complete (§C. elegans unified P3-P6 run, §Controlled intervention evaluation); next is held-out-target intervention validation, then ZAPBench/Allen adapters |
| HIGH | Cluster bring-up | Multi-node smoke (torchrun --nnodes=2 --nproc_per_node=4, DeepSpeed ZeRO-2/3): forward/backward, checkpoint save/load, elastic resume; profile 16/24/32/48–80 GiB tiers and re-measure target-topology VRAM and throughput before any full-scale run |
| MED | Phase -1 Magi v2 pretraining | configs/data/eegdenoisenet.json now provides an EEG-only clean/artifact loader; materialize the public arrays, then run topology-specific checkpoint/resume, masked-leakage, causal-NTP, and EMA validation |
| MED | Time constants | latent_dt single-source done (species sample rate → model; OU integrator synced, default-preserving). Remaining: wire-or-drop SlowGateTransition (zero callers today); document seconds conventions for monitors/free-run cadence |
| MED | Marginal recalibration | Per-channel raw-unit gain is anchored at the reconstruction boundary; full OU/TDE-RICA component-level marginal recalibration on held-out rollouts remains before claiming full output-marginal preservation |
| MED | Subject adaptation | Re-review the low-rank test-time adaptation design and wire into Stage 3 |
The first balanced remote optimization run used the unified C. elegans
profile with latent_velocity_scale=0.02 and a normalized calcium
reconstruction mixture:
{"correlation": 0.55, "corr_diff": 0.20, "wasserstein1": 0.25}This follows the gKDR-GMM reference priorities: preserve variance/marginal scale, retain temporal transitions, and reduce distributional discrepancy without discarding correlation structure. The evaluated checkpoint was:
Stage 1 P1, step 5000
Remote validation used eight subject-held-out samples, three autonomous rollout
windows, and the neighboring ../TDE-RICA implementation. Results:
| Metric | Mean |
|---|---|
| Native correlation-matrix MSE | 0.3423 |
| Native autocorrelation MSE | 0.2629 |
| Native variance log-RMSE | 1.6432 |
| TDE-RICA W1 | 0.7366 |
| TDE-RICA KL mean | 13.3837 |
| TDE-RICA kernel transition distance | 2.4715 |
| Real occurrence lag-1 | 0.9853 |
| Generated occurrence lag-1 | 0.9048 |
The checkpoint improves dramatically over the prior step-5000 artifact, whose variance log-RMSE was 12.5475 and TDE-RICA W1 was 49.2513. The user-specified reference gates are W1 <= 0.05, mean KL < 5, kernel-transition distance < 1, and raw standard-deviation maintenance near 1. The optimized checkpoint's measured values are:
| Gate | Target | Optimized mean | Status |
|---|---|---|---|
| Raw std ratio | near 1 | 0.8231 | improved, not near enough |
| Native variance log-RMSE | lower | 1.6855 | improved, still high |
| TDE-RICA W1 | <= 0.05 | 0.7366 | fail |
| TDE-RICA KL mean | < 5 | 13.3837 | fail |
| Kernel transition distance | < 1 | 2.4715 | fail |
| Generated occurrence lag-1 | near real 0.9853 | 0.9074 | fail |
The W1-heavy interrupted experiment at step 3323 improved statistical
correlation (0.2039) and tail-variance stationarity (0.9808) but worsened
mean raw standard-deviation maintenance (0.6251), W1 (0.8398), KL
(17.6089), and occurrence (0.8214). It was not selected.
The current artifact is therefore a valid optimization checkpoint and a material improvement, but it does not meet the requested gKDR-GMM-like thresholds. More training with the present objective should not be claimed to solve the gap; the next needed experiment is a calibrated distributional objective in the same TDE-RICA component space, with a fixed reference basis and scale calibration rather than raw calcium reconstruction alone.
Post-training verification passed in the remote environment: the checkpoint
loaded with no missing or unexpected keys, forward_modalities() emitted
finite (B,3,C,T) reconstructions and (B,3,192) latent rollouts, the mixed
loss returned a finite scalar, the threshold-aware evaluator ran successfully,
and tools/real_data_sanity.py compiled/imported. The local workstation
environment does not provide NumPy, so numerical tests were run remotely.
The unified three-source manifest (1,056 rows: 24 Toyoshima salt, 113 Randi PumpProbe, 919 HuggingFace) was trained end-to-end for the first time. The run continues Stage 1 at P3 from the previous P3 artifact and trains P3-P6 at 3,000 steps per phase on the unified corpus:
python -u train.py \
--config configs/species/c_elegans_unified.json \
--data configs/data/c_elegans_unified.json \
--deepspeed_config configs/ds_config_zero2.json \
--phase 1 --phase_steps 3000 \
--resume "<previous>/best_Stage 1 P3: Multi-Horizon Forecasting_step2000.pt" \
--log_dir ../logs_unified_p3p6 --checkpoint_dir ../ckpt_unified_p3p6Measured: 950 training rows / 106 validation rows (subject-grouped, seed 0),
1,504 s wall clock for 10,000 loop steps on one RTX 3080 Ti at batch 1
(6.6 steps/s: P3 1,000 steps in 2.50 min, P4 3,000 in 7.43 min, P5 3,000 in
7.62 min, P6 3,000 in 7.52 min; peak GPU allocation 0.4 GiB). The logged total
loss falls from 0.0515 at P4 step 380 to 0.0067 at P6 step 2800. Four phase-end
checkpoints were written
(checkpoint_Stage 1 P3/P4/P5/P6..._step3000.pt).
Four defects blocked or silently invalidated this run and were fixed in the same pass:
| Defect | Effect before the fix | Fix |
|---|---|---|
configs/data/c_elegans_*.json roots used ../../data_ladder/... |
Resolved to <repo>/data_ladder, which does not exist, so every --data configs/data/... invocation died in the loader probe |
Roots are ../../../data_ladder/..., matching the sibling-layout convention already used by configs/data/eegdenoisenet.json |
| Phase budgets were fixed at full scale (P4 50k, P5 100k) | A bounded corpus could only be run by signalling a preemption checkpoint | --phase_steps N caps each selected phase and clamps its warmup |
train_phase() saved only on the 5,000-step cadence |
A 3,000-step phase finished and persisted nothing; the completed run would have produced zero usable checkpoints | train_phase() saves at phase end, matching the Magi pretraining path |
| Evaluator had no source or horizon breakdown | Per-source questions could not be answered from one validation run | tools/remote_free_run_eval.py labels each sample's ladder source and reports per-horizon native metrics, aggregate_by_source, and aggregate_by_horizon; unscoreable windows are recorded as skips instead of aborting an audit |
Checkpoint Stage 1 P6: Long-Context Expansion, step 3000 (loaded with zero
missing and zero unexpected keys) evaluated on the 106-row held-out validation
split, three autonomous rollout windows per sample, TDE-RICA fitted per
sample's context window, --no-d3:
| Group | n | Raw std ratio | Corr-matrix MSE | Autocorr MSE | Variance log-RMSE | TDE-RICA W1 | TDE-RICA KL | Gen. lag-1 | Real lag-1 |
|---|---|---|---|---|---|---|---|---|---|
| All | 106 | 0.6306 | 0.3063 | 0.2735 | 1.4686 | 0.8366 | 14.72 | 0.8724 | 0.9799 |
| HuggingFace (no control) | 94 | 0.6042 | 0.3247 | 0.2805 | 1.4150 | 0.8294 | 14.24 | 0.8759 | 0.9873 |
| Randi PumpProbe | 11 | 0.8621 | 0.1645 | 0.2263 | 1.9455 | 0.9140 | 18.51 | 0.8347 | 0.9153 |
| Toyoshima salt | 1 | 0.5686 | 0.1391 | 0.1349 | 1.2556 | 0.6638 | 17.63 | 0.9573 | 0.9900 |
Persistence baseline on the same rows scores 0.4936 correlation-matrix MSE against the model's 0.3063, so the free run beats "hold the last observation".
Same split, same evaluator, earlier checkpoint of the same run (P3, step 3000) versus the final P6 checkpoint:
| Metric | P3 step 3000 | P6 step 3000 |
|---|---|---|
| Raw std ratio | 0.5351 | 0.6306 |
| Corr-matrix MSE | 0.3145 | 0.3063 |
| Autocorr MSE | 0.2877 | 0.2735 |
| Variance log-RMSE | 1.7423 | 1.4686 |
| TDE-RICA W1 | 0.8671 | 0.8366 |
| TDE-RICA KL | 17.29 | 14.72 |
| Generated occurrence lag-1 | 0.7733 | 0.8724 |
P6 dominates P3 on every axis measured, which is the first within-run evidence that the P4-P6 phases (routing, physics-inspired constraints, long-context continuation) improve the free-run distribution rather than only the reconstruction loss.
Horizon breakdown on the held-out split (each row's (K,C,T) targets are three
successive 15-second windows, so horizon k spans seconds [15k, 15(k+1))
after the context):
| Horizon | n | Raw std ratio | Corr-matrix MSE | Autocorr MSE | Variance log-RMSE |
|---|---|---|---|---|---|
| 1 | 106 | 1.129 | 0.3278 | 0.2176 | 2.0104 |
| 2 | 106 | 0.882 | 0.4160 | 0.2193 | 1.8951 |
| 3 | 106 | 1.055 | 0.3619 | 0.2057 | 1.8944 |
Scale and autocorrelation hold across the three horizons; correlation structure degrades in the middle window rather than monotonically.
Source audits use the source-scoped profiles
(configs/data/c_elegans_unified_audit_{toyoshima,randi,hf}.json, same
15-second window and 282-wide control contract as training, val_frac=1.0 so
every row of that source is scored). These are not held-out splits: the
Toyoshima and Randi audits include rows used for training, and the HF audit
scores 918 of the 919 HuggingFace rows, so they measure the corpus fit, not
generalisation. They exist because the mixed validation split carries only one
salt row:
| Source audit | n scored | Raw std ratio | Corr-matrix MSE | Autocorr MSE | Variance log-RMSE | TDE-RICA W1 | TDE-RICA KL | Gen. lag-1 | Real lag-1 |
|---|---|---|---|---|---|---|---|---|---|
| Toyoshima salt | 23 | 0.7002 | 0.1164 | 0.1163 | 1.0061 | 0.7512 | 17.27 | 0.9215 | 0.9778 |
| Randi PumpProbe | 112 | 0.9447 | 0.1568 | 0.2141 | 1.6342 | 0.9713 | 19.52 | 0.8204 | 0.9178 |
| HuggingFace | 889 (29 skipped) | 0.6232 | 0.3157 | 0.2761 | 1.3854 | 0.8781 | 13.88 | 0.8812 | 0.9869 |
The 29 HuggingFace rows are skipped, not dropped silently: each has fewer than
two calcium channels whose entire future window is observed, and the report
records them under skipped.
Against the gKDR-GMM reference gates (W1 ≤ 0.05, mean KL < 5, kernel
transition < 1, raw standard deviation near 1) the run is still short: W1 is
0.84 on the held-out split, KL is 14.7, and the raw standard-deviation ratio is
0.63. The kernel-transition gate is not evaluated here (--no-d3); it stays
open until a stratified D3 run reproduces it. The four P6 free-run reports
above (106 + 23 + 112 + 918 rows) completed in ≈11 minutes of wall clock on the
same single card, and the P3 comparison report on the 106-row split took
≈6 minutes.
Free-run scoring says nothing about whether the recorded drive actually steers
the model, because it deliberately runs without control. tools/intervention_eval.py
rolls the same context three times — with the recorded drive, with an explicit
all-zero drive, and with no drive argument — and compares the recorded-versus-
zero difference against the held-out future's own deviation from persistence:
python tools/intervention_eval.py \
--config configs/species/c_elegans_unified.json \
--data configs/data/c_elegans_unified_intervention_randi.json \
--checkpoint "<P6>/checkpoint_Stage 1 P6: Long-Context Expansion_step3000.pt" \
--output intervention_randi.json --split val --max-samples 112 \
--rollout 3 --device cudaTwo evaluation-construction facts had to be measured first:
- The val split's first eligible window is pre-stimulus for both
intervention sources (measured first-drive onsets: Randi 19–126 s, median
24 s; Toyoshima salt 30–630 s, median 630 s). Every control-bearing row
therefore scored
drive_norm = 0. Thecontrol_alignprofile key starts each evaluation window at the first frame whose control value is non-zero; with it, 23/23 salt and 110/112 Randi rows carry an active drive. - The documented "zero stimulus is an exact no-op" invariant did not hold
for a trained checkpoint. Both the readout (
perturbation_map) and the input-conditioned gate (gate_map) are only zero-initialised; their trained biases keep responding to the mere presence of a control tensor (measured 9–13 × 10⁻³ maximum absolute calcium deviation against the no-control rollout). Both paths are now masked by drive presence, which makes the invariant structural rather than initialisation-dependent: re-running the same evaluation reportszero_drive_max_abs_diff = 0.0exactly for every row.
The P6 checkpoint reported below predates a contract fix in
training_loop.partition_generic_batch: when *_future controls are present,
each full future control window is now reduced independently for its matching
rollout horizon; the old path reduced the current context window into rollout
segments. The historical intervention numbers remain valid for that checkpoint,
but a post-fix checkpoint is required before deciding whether a dedicated
intervention loss is necessary.
Results on the P6 checkpoint, drive-aligned windows, three horizons:
| Arm | Rows | Drive-active | Model effect ‖·‖ | Real effect ‖·‖ | Gain | Direction cosine | Channel-pattern corr |
|---|---|---|---|---|---|---|---|
| Toyoshima salt (stimulus vs no-stimulus) | 23 | 23 | 0.507 | 147.99 | 0.0036 | 0.001 | −0.019 |
| Randi PumpProbe (optogenetic vs no-op) | 112 | 110 | 1.446 | 87.15 | 0.0179 | −0.001 | 0.007 |
Targeted-versus-non-target channels, Randi rows with an active target gate
(106 rows). Gate k addresses calcium channel k − 1 in the recording-local
channel order:
| Quantity | Model | Data |
|---|---|---|
| Targeted / non-target mean |effect| ratio | 1.011 | 1.676 |
Interpretation, stated at the strength the measurement supports:
- The conditioning pathway is materially inert in this checkpoint: the drive moves the free run by 0.4 % (salt) to 1.8 % (Randi) of the distance the recording actually moved, with no directional agreement (cosine ≈ 0) and no spatial agreement across channels.
- The data itself carries the expected specificity — Randi's targeted channels respond 1.68× more than non-targeted ones — and the model reproduces none of it (ratio 1.01). The intervention signal is present in the corpus and absent from the model's response.
- This does not invalidate the free-run numbers above: those rollouts pass
perturbation=None, so they never enter the masked path. It does mean the parameter/objective question for the next run is the control pathway's supervision, not just the reconstruction mixture. perturbation_mapstarts at exactly zero, so every step of control learning is driven purely by forecast gradients. Ten thousand unified steps are evidently not enough to grow that path; the next experiment should either raise the control pathway's learning signal (dedicated intervention loss, larger weight on driven windows) or verify that the latent dynamics can be steered at all before scaling the corpus.
Every free-run claim in the previous two sections rests on this evaluation
matrix, all of it on the same checkpoint (Stage 1 P6: Long-Context Expansion, step 3000) and one RTX 3080 Ti:
| Report | Rows scored | Windows | TDE-RICA | Artifact |
|---|---|---|---|---|
| Held-out validation split | 106 | 3 × 15 s | per-sample context fit | eval_unified_p6_val_full.json |
| Toyoshima salt audit | 23 | 3 × 15 s | per-sample fit | eval_unified_p6_salt_arrays/ + eval_unified_p6_audit_toyoshima.json |
| Randi PumpProbe audit | 112 | 3 × 15 s | per-sample fit | eval_unified_p6_audit_randi.json |
| HuggingFace audit | 889 (29 skipped) | 3 × 15 s | per-sample fit | eval_unified_p6_audit_hf.json |
Dynamics subset (--full-d3) |
8 | 3 × 15 s | D3 estimators enabled | eval_unified_p6_val_d3.json |
| Intervention arms | 23 + 112 | 3 × 15 s | n/a | intervention_p6_*_aligned.json |
Per-sample and per-channel distributions are saved from the --save-arrays
matrices by tools/per_channel_distributions.py
(per_channel_distributions.json + .npz):
| Arrays | Samples | Active channels/sample | Generated/real channel σ ratio (mean, q05, q95) | Lag-1 real → generated |
|---|---|---|---|---|
| Validation split (P6) | 106 | 47.4 | 0.954, 0.253, 2.372 | 0.967 → 0.225 |
| Validation split (P3) | 106 | 47.4 | 0.902, 0.235, 2.367 | 0.967 → 0.064 |
| HuggingFace audit | 889 | 36.9 | 0.918, 0.236, 2.411 | 0.976 → 0.256 |
| Randi audit | 112 | 114.0 | 1.406, 0.422, 2.746 | 0.894 → 0.060 |
| Toyoshima salt audit | 23 | 67.1 | 0.807, 0.416, 1.570 | 0.750 → 0.324 |
Two readings: P6 roughly triples the generated lag-1 of P3 (0.064 → 0.225) without changing the scale distribution, and the Randi arm is the weak one — over-dispersed per channel (1.41) with almost no temporal correlation generated (0.060 against a real 0.894). The ratio excludes channels whose real trace is near-constant (below 1e-3 of the sample's median σ); including them makes the pooled mean meaningless (measured 5.7 × 10⁸ before the guard).
The full D3 run closes the kernel-transition gate question on an 8-row subset (dynamics estimators are the expensive part):
| Metric | 8-row subset |
|---|---|
| Raw std ratio | 0.930 |
| Corr-matrix MSE | 0.242 |
| TDE-RICA W1 | 1.498 |
| TDE-RICA KL | 18.27 |
| Kernel transition | 5.259 |
| Jacobian distance | 8.767 |
| Expansion / rotation diff | 0.211 / 0.337 |
| Transfer entropy real→gen / gen→real | 4.094 / 2.472 |
| Occurrence lag-1 (real → generated) | 0.981 → 0.909 |
The subset's other metrics differ from the 106-row means (it is a harder
draw), so only the gate conclusion is drawn from it: kernel transition is
5.26, more than 5× above the requested < 1, and it is not re-measured here.
Independent-seed check. The same schedule (P4-P6 at 3,000 steps per phase, from
the same P3 artifact) was re-run with --seed 1 and evaluated on the same
106-row held-out split:
| Metric | Seed 42 | Seed 1 |
|---|---|---|
| Raw std ratio | 0.6306 | 0.5914 |
| Corr-matrix MSE | 0.3063 | 0.2924 |
| Autocorr MSE | 0.2735 | 0.2776 |
| Variance log-RMSE | 1.4686 | 1.5242 |
| TDE-RICA W1 | 0.8366 | 0.8402 |
| TDE-RICA KL | 14.72 | 15.41 |
| Generated occurrence lag-1 | 0.8724 | 0.8687 |
Every axis agrees within a few percent, so the reported levels are not an artefact of one initialization. An earlier seed-1 run at half the per-phase budget (1,500 steps) landed in the same range (std ratio 0.527, W1 0.861, KL 16.98), which also bounds how much of the P3 → P6 gain is step budget.
Checkpoint/profile pairing is now explicit rather than assumed: train.py
hashes the resolved species profile into every checkpoint
(config_sha256), and the evaluators compare it against the --config they
are given, reporting config_match and printing a warning before loading
weights when the pair differs:
[config] WARNING: checkpoint was produced by a different profile
(.../c_elegans_unified.json); evaluating (.../c_elegans.json) anyway
Verified both ways: the same profile reports config_match: true with
identical hashes; a different profile warns first and then fails on the
shape mismatch it would have hit silently otherwise. Checkpoints written
before this change report config_match: null and load as before.
What the measurements above justify changing, separated from what they do not:
| Finding (measured) | Reading | Next action |
|---|---|---|
| P6 dominates P3 on every free-run axis (W1 0.867 → 0.837, KL 17.3 → 14.7, occurrence lag-1 0.773 → 0.872, std ratio 0.535 → 0.631) | The P4-P6 phases (routing, physics-inspired regularization, long-context continuation) improve the distribution, not just the reconstruction loss | Keep the phase structure; do not re-tune it before the objective changes |
| Held-out W1 0.837 / KL 14.72 vs gates 0.05 / 5; kernel transition 5.26 vs < 1 | A 25-minute, 10k-step run is not near the gates, and the P3 → P6 gain is modest, so more steps of the same objective are not the answer | Add an explicit distributional term in the fixed TDE-RICA component space (the gKDR-GMM reference basis), as concluded before this run |
| Drive effect ≈ 0.4 % (salt) / 1.8 % (Randi) of the real effect; cosine ≈ 0; targeted ratio 1.011 vs data's 1.676 | The control path is not merely weak, it is untrained: perturbation_map starts at exactly zero and only forecast gradients can grow it |
Supervise the control path directly (intervention-matched loss, or drive-weighted sampling) and re-measure with tools/intervention_eval.py before any steerability claim |
| Randi arm: per-channel σ ratio 1.406 with generated lag-1 0.060 (real 0.894) | The latent dynamics are too fast for the 2 Hz Randi windows; the HF arm (0.918 / 0.256) behaves far better on the same checkpoint | Investigate source-conditioned dynamics (per-source adapter or mixing weight) or horizon supervision on Randi before pooling sources in one objective |
| Aggregate held-out std ratio 0.631 vs per-channel mean 0.954 | Contraction in the aggregate is partly a channel-mix effect (large-σ channels shrink more), not a uniform gain error | Report both numbers; treat the per-channel distribution, not the pooled ratio, as the scale gate |
Nothing above is a claim that the architecture is wrong. The evidence supports one objective change (distributional term), one supervision gap (control pathway), and one data-mix question (Randi), all of which are measurable with the tools added here.
The gKDR-GMM comparison must use the Toyoshima cohort and its fixed basis,
not the unified HF-first federated validation path. Octave inventory of
test20220308_tderica_expansion_visualize_results_2_captured_comp_and_coeff.mat
confirmed:
coeffEmbed2 14 × 300 × 177
compEmbed2 5701 × 24 × 14
tcrsN3 6000 × 177 × 24
strNamesOrdered 177 × 1
Export or pass this MAT file to tools/tderica_free_run.py with
--basis. The tool projects both signals through the same coeffEmbed2
basis and records fit_on_evaluation_window=false. This matches the
Toyoshima/gKDR settings: 24 samples, 177 ordered neurons, 300-frame delay
embedding, time_step=5, embed_step=10, and the canonical-name plus
lag-20 autocorrelation QC.
python tools/tderica_free_run.py \
--real toyoshima_real.npy \
--generated toyoshima_generated.npy \
--basis test20220308_tderica_expansion_visualize_results_2_captured_comp_and_coeff.mat \
--dt-s 0.2426317083 \
--tderica ../TDE-RICA \
--output toyoshima_fixed_basis_report.jsonFor a baseline-comparable report, pass a persistence trace from the same
evaluation window with --persistence. The real, generated, and persistence
signals are then projected through the same fixed basis; the JSON keeps the
baseline report under baseline_reports.persistence and records model-minus-
baseline metric deltas under baseline_comparison. The evaluator validates
finite coefficients and unique ordered channel names and fails closed when the
basis artifact is absent or incompatible. For traces with missing channels, it
intersects finite delay-embedded channel masks across real, generated, and
persistence inputs and records the retained canonical names in the report. It
does not fit a replacement basis on the evaluation window.
The Pareto utilities in diagnostics/free_run_metrics.py treat raw standard
deviation error, W1, KL, kernel-transition distance, occurrence gap, and native
correlation error as separate minimization axes. They must not be collapsed
into a biologically interpretable scalar without a declared normalization.
The unified C. elegans root is a federated training corpus. It is useful for training but is not the reference space for the published gKDR-GMM numbers. TDE-RICA/gKDR components are defined by the Toyoshima cohort and its 177-neuron ordering. Use the fixed-basis tool against Toyoshima signals when comparing W1, KL, kernel transition, or occurrence values to that reference. Do not fit a basis on the evaluation window.
The captured reference MAT file was inventoried with Octave and contains
coeffEmbed2 (14,300,177), compEmbed2 (5701,24,14), tcrsN3 (6000,177,24), and strNamesOrdered (177,1). The reviewed manifest is
configs/data/c_elegans_toyoshima_fixed_basis.json.
diagnostics.free_run_metrics.pareto_front() and
tools.tderica_free_run.select_pareto_checkpoint() keep variance, W1, KL,
kernel transition, occurrence, and native statistical errors as separate
axes. A checkpoint is a candidate only when its report was produced in the
same fixed basis and all required metrics are finite. The compromise selector
is a convenience for an explicit archive, not a substitute for declaring
metric normalization and biological acceptance gates.
Reproducibility command (remote paths must be adapted). The full 6000-frame Toyoshima Fréchet calculation is computationally expensive; use a bounded evaluation window for routine smoke tests, and reserve the full trace for the final evidence run:
PYTHONPATH=/path/to/repo:/path/to/TDE-RICA/TDE-RICA \
python tools/remote_free_run_eval.py \
--config configs/species/c_elegans_unified.json \
--data /path/to/c_elegans_opt_profile.json \
--checkpoint "/path/to/checkpoint_Stage 1 P1_step5000.pt" \
--output /path/to/eval/full_5000.json \
--split val --max-samples 8 --rollout 3 --device cuda \
--tderica /path/to/TDE-RICA/TDE-RICA| LOW | Housekeeping | Move __main__ demos to examples/ (tests/conftest.py already bootstraps imports for pytest) |
| MED | Unwired config keys | replay_species and the implicit cross-species mixing contract were removed (2026-09): each ladder stage trains its own model against its own ladder, so the flag had no consumer. Remaining inert key: human.json lists behavior with no declared role |
| MED | MoE structure scope | With use_generic_moe=False (default) experts add an unconstrained velocity, so the composite field is not GENERIC (measured: moe_velocity_share ≈ 0.76, moe_energy_alignment ≈ 0.16 at init). Ablations in §6 therefore test the backbone. Structured mode is affordable now (96M vs 26M, was 379M) but needs a training comparison before becoming the default |
| MED | Trajectory discriminator supervision | The CFTS TrajectoryDiscriminator that selects paths is shaped only indirectly (soft weights → endpoint EFE → ActionLoss, weight 0.05 in Stage 3). "Best trajectory" is closer to a learned preference direction than a supervised selector |
| MED | MT-KDA scope | The multi-time-scale filter is applied inside VelocityBrain, so it smooths the backbone but not the MoE velocity added afterwards in _latent_step. Its states are also detached (no BPTT through the filter) and reset per forward |
| LOW | Perception correction | Now a free-energy gradient step (active_inference_step, default 0.1) instead of a detached pull toward the rollout endpoint. Needs a step-size sweep on real runs |
| LOW | Perturbation validation | Control input is wired from stimulus-role data (roles → perturbation reduction → conditioned gates, --noise_mode policy); optogenetic target-gated encoding and pulse-preserving reduction are now implemented. Remaining: acquire Randi/zebrafish/mouse intervention corpora, register neuron versus population target scopes and fixed vocabularies, use peak for sparse pulses, and run held-out-target validation |
| MED | Conditioned diffusion (P2.2) | Diffusion gain σ(u) conditioned on control (zero-init bounded gain, default identity) — deferred until control metrics are reviewed on salt/HF ingests |
| RESOLVED | torch_xla rebinds nn.GRU globally | Cause confirmed: importing torch_xla rebinds the class itself (nn.GRU -> torch_xla.experimental.gru.GRU) via torch_xla/_patched_functions.py:69, nn.GRU = _pathch_module(nn.GRU, ScanGRU) - a global rebind with no env guard, so CPU/CUDA input raises RuntimeError: Expected all tensors in the given list to be XLA tensors. Minimal repro: the same call works with torch_xla unimportable and fails the moment it is imported. runtime.device_utils.device_aware_gru() subclasses the installed class and routes per call from the input tensor: XLA input keeps the scan path, anything else uses the native class that _pathch_module stores as _orig. Hosts without torch_xla get plain nn.GRU untouched, so CUDA/CPU code paths are unchanged, and parameter names stay identical (checkpoint-compatible) |
| MED | Multi-worker DataLoaders | build_data_loaders defaults to num_workers: 4, so any profile omitting the key forks workers. Two independent hazards: (1) fork-after-threads deadlocks on runtimes that start threads (torch_musa does) - observed hanging the suite on the MTT S4000, waiting forever in _try_get_data; (2) on that same MUSA stack multi-worker loading cannot work at all, because torch_musa's patched torch/multiprocessing/reductions.py tests storage.is_musa, which the installed torch build does not define, so any tensor crossing a worker boundary raises AttributeError and the parent hangs. Reproduced with a trivial dataset, so it is not this repo's code. spawn was tried as a fix and reverted: it imposes an __main__ guard on every caller and does not help where (2) applies. Test profiles pin num_workers: 0; use in-process loading on MUSA until the torch/torch_musa versions match |
| MED | TODO: re-verify GRU fix on TPU | The cause was verified on a Colab TPU v5e1, and an earlier always-native variant took the suite from 3 failed, 138 passed to 144 passed, 2 skipped. The implementation has since changed to per-call device dispatch, which that run does not cover - it is verified locally only (12 tests). Re-run the TPU suite and additionally confirm a forward on a real XLA tensor keeps using the scan implementation. Blocked on Colab TPU capacity being unavailable |
| LOW | Deferred features | See table below |
| Item | Trigger to revisit |
|---|---|
| EEG decoder exact 2-layer spec | Reconstruction quality poor |
| Hard modality activation in router | Soft blend causes gradient conflicts |
| Factorized attention (spatial/temporal heads) | Memory-constrained scenarios |
| fMRI functional-connectivity input mode | Alternative to ROI mode needed |
| Virtual montage mapping | BIOT fallback insufficient |
| OU noise temperature scan | Critical-dynamics analysis requires it |
| MNN moment embedding consistency loss | Research-gated (statistical mechanics) |
tsallis latent sparsity |
Sign bug rewards concentration - fix before any enable |
velocity_smooth temporal TV |
Implemented (2026-09): true temporal TV over delta_z_sequence, which the model emits for num_steps > 1 rollouts. Still weight 0 by default; enable via phase rollout_steps > 1 + weight |
cross_soft hub contrastive |
Enable only when the trainer supplies cross_modal_labels |
- No thermodynamic claims: energy/entropy are learned fields; L(z) antisymmetry is not a Poisson structure (Jacobi stays an offline diagnostic); EPR/signed-work/divergence terms are named proxies and are not rollback triggers unless explicitly supplied.
- No universal criticality claim: homeostat and monitors are control-proxy regulation.
- Spectral losses apply to reconstructed signals only; 1/f lives in one
place; distribution and criteria choices per modality follow the
TDE-RICA/C. elegans validation practice (see
SPECIES_PROFILES). - Convention: torch-native code; numpy only for offline diagnostics.
- Conditioning semantics: control readouts are bias-free with a zero-initialized readout, and the input-conditioned gates are masked by drive presence, so an all-zero stimulus is an exact no-op in any trained state; measured by
tools/intervention_eval.pyaszero_drive_invariant_max_abs_diff. Noise is off unlessnoise_modeexplicitly enables it. - Species/channel metadata is conditioning, not homology.
- Verified by inspection (2026-09):
EngramLandscape.consolidate()is an unimplemented future API with zero callers (never on a hot path); the MoE lesionelse: passbranch is intentional drop semantics for lesioned routed experts.
Code: development checkout in Google Drive (マイドライブ/a/brain_moe_pinn,
self-contained, git-tracked); canonical HPC path
/home/yanlu/Documents/a/brain_moe_pinn/.
Fundamental principle: Model velocity $\Delta z$, not absolute position.
Irreversibility is a feature, not a bug.