PaperlayerView exact revision on arXiv

Source-linked AI summary

Robust Neural Stimulation Response Modeling Through Meta-Learning and Pretraining

Matthew J Bryan, Daniel C Muir, Felix Schwock, Azadeh Yazdan-Shahmorad, Rajesh P N Rao

arXiv:2608.26649v1cs.LG

TL;DR

Model-based closed-loop stimulation is constrained by catastrophic forecast failures and calibration demands that may not fit clinical sessions. The paper extends TBFMs with cross-session pretraining and MAML, evaluating the approach on 40 optogenetic sessions from two rhesus macaques. Pretrained models improve robustness and reduce calibration requirements while retaining practical inference characteristics.

  • Problem

    Closed-loop stimulation requires forecasting models that are robust to catastrophic session failures and efficient enough to calibrate within clinical constraints.

  • Method

    The paper extends TBFMs with a cross-session architecture, MAML pretraining, and test-time adaptation for unseen stimulation sessions.

  • Results

    Pretrained TBFMs significantly exceed single-session models in sample efficiency and robustness, with MAML providing an additional measurable improvement over co-adaptive pretraining.

  • Takeaways & Limitations

    Cross-session structure in stimulation responses can support pretraining that produces practical gains in neural stimulation efficiency and reliability.

  • Takeaways & Limitations

    Findings are limited by data from two rhesus macaques, S1/M1 regions, and one optogenetic stimulation modality.

Abstract

from arXiv · show

Objective: Model-based closed-loop neural stimulation holds promise for therapeutic applications ranging from Parkinson's disease to sensory restoration, but deployment has been limited by two obstacles: 1) forecasting models for predicting the consequences of stimulation fail catastrophically on a meaningful fraction of sessions, and 2) per-session calibration requirements are often incompatible with clinical constraints. We address both by demonstrating, for the first time, that meta-learning and pretraining can be applied to neural stimulation response modeling. Methods: Temporal basis function models (TBFMs) forecast state-dependent neural responses to stimulation. We extend TBFMs with cross-session pretraining using a novel architecture and algorithm based on model-agnostic meta-learning (MAML), evaluating them on 40 sessions of optogenetic stimulation in primary sensorimotor cortex of two non-human primates. Results: Meta-learning substantially reduces catastrophic forecast failure: for a 1k calibration set size, sessions with test R-squared < 0.05 drop from 16 of 40 (single-session training) to 1 (MAML-pretrained), and prediction intervals become significantly narrower (p < 0.05). Calibration requirements are reduced by 50-90% at matched accuracy, enabling experiments otherwise infeasible within clinical session-time constraints. Conclusion: Our results demonstrate that cross-session structure in stimulation responses is consistent enough to support pretraining, providing the first empirical evidence that meta-learning approaches are viable for neural stimulation. Significance: The robustness and sample efficiency gains directly address known obstacles to deploying model-based stimulation controllers. Our results motivate community efforts to assemble standardized multi-site stimulation datasets and to further explore meta-learning for robust closed-loop stimulation.

I. INTRODUCTION

Closed-loop neural stimulation can adapt to brain state but requires robust, sample-efficient, low-latency forecasting. This work extends TBFMs with cross-session pretraining and meta-learning to reduce catastrophic failures and calibration demands.

  • Closed-loop stimulation adapts treatment to neural or behavioral state but requires robust, sample-efficient, accurate, and low-latency controllers.
  • TBFMs provide single-trial, multi-step, spatiotemporal forecasts of state-dependent stimulation responses with high sample efficiency and low latency.
  • Pretraining past sessions with MAML substantially improves forecasting robustness and reduces catastrophic accuracy degradation on noisy or nuisance-variable sessions.
  • Pretrained TBFMs adapt to future sessions with calibration datasets 50−90% smaller than single-session TBFMs, enabling experiments otherwise infeasible within session-time constraints.
  • The approach combines a cross-session architecture, MAML pretraining, test-time adaptation, and empirical gains in robustness, calibration efficiency, and forecast accuracy.
  • The evaluation uses 40 optogenetic stimulation sessions from two rhesus macaques while recording cortical activity with a 96-channel µECoG grid.

B. Introduction to temporal basis function models (TBFMs)

TBFMs forecast stimulation-evoked neural activity from a short pre-stimulation brain-state window and stimulation parameters. Their design targets sample efficiency, training efficiency, and real-time inference while accommodating session-specific stimulation variation.

  • Key design goals:: TBFMs address modest-data, computational, and latency constraints relevant to closed-loop experiments and real-time control.
  • Key design goals:: TBFMs use a short neural-data runway and stimulation parameters to generate single-trial, multi-step forecasts of stimulation responses.
  • Key design goals:: The runway contains 20ms of neural data, followed by an assumed 20ms delay before stimulation delivery to reflect worst-case control-loop latency.
  • Key design goals:: Within-session stimulation parameters are fixed, whereas inter-pulse intervals and target locations vary across sessions and must be handled during pretrained-model adaptation.
  • Architecture:: The model combines stimulation descriptors, shared basis vectors, normalized runway data, and an affine weight estimator to produce channel-wise forecasts.
  • Architecture:: The forecast for each channel is a weighted sum of shared basis vectors, with weights estimated from the normalized runway.

3) Training:

Training extends TBFMs from single-session forecasting toward cross-session models, while preserving state-dependent prediction and practical inference efficiency.

  • Training: Training uses early-session trials and evaluates on later trials to test robustness to within-session drift.Mini-batches contain 5k training trials, while evaluation uses the final 2.5k trials.
  • Training: TBFMs forecast stimulation responses conditioned on brain-state-dependent initial conditions.A high initial-state bin shows a significantly different response from the others in a single-channel, single-session example.
  • Training: Compilation precomputes basis sets for discrete stimulation parameters, reducing typical inference processing by approximately 34%.The compiled basis dictionary avoids recomputing bases during inference.
  • Training: TBFMs previously achieved the highest overall accuracy, two-orders-of-magnitude lower inference latency, and better sample efficiency than the compared LSSM and LSTM models.The comparison found no efficiency–accuracy trade-off on those data.
  • Training: The cross-session architecture adds per-session normalization, latent-space autoencoders, a shared TBFM, and normalized basis weights.Median/IQR normalization and session-specific autoencoders address small calibration sets and differing channel availability.

1) Per-session resting state context:

The model uses resting-state and stimulation-specific session context to shape shared temporal bases while adapting to session-specific response characteristics.

  • Per-session resting state context: Baseline autocorrelation statistics help shape temporal bases because they correlate with stimulation-forecast predictability.Lower autocorrelation motivates more conservative forecasts that revert quickly toward the channel mean.
  • Per-session resting state context: Resting-state context summarizes channel-wise baseline A-ACFs using their 25th, 50th, and 75th percentiles before stimulation begins.The resulting context vector is computed from baseline data in advance.
  • Per-session resting state context: A learned stimulation context is optimized during training and test-time adaptation to customize bases for each session.This reflects shared response structure alongside unique session-specific characteristics.
  • Per-session resting state context: An MLP converts stimulation context into a low-rank residual correction that is added to the shared basis matrix.The reported residual rank is 16, enabling adaptation with few test-time parameters.
  • Per-session resting state context: MAML simulates calibration during training by adapting a stimulation context on a support set and updating shared parameters using a query set.Random support-set sampling trains adaptation across possible calibration subsets.

4) Test-time adaptation (TTA) procedure:

Test-time adaptation specializes the pretrained cross-session model to an unseen session using a small calibration set, after which compilation simplifies inference.

  • Test-time adaptation (TTA) procedure: TTA adapts a pretrained TBFM to forecast stimulation effects for an unseen session using a smaller calibration set than single-session training requires.The procedure is the main mechanism for transferring pretraining to new sessions.
  • Test-time adaptation (TTA) procedure: The unseen session’s TTA support set consists of newly collected calibration data rather than a random split from existing training data.This distinguishes deployment-time adaptation from the simulated support sets used during MAML training.
  • Test-time adaptation (TTA) procedure: During TTA, the session-specific autoencoder is fully trained, the stimulation context is optimized, and the basis generator is fine-tuned at learning rate 10^-6.The autoencoder is warm-started with PCA on the calibration data.
  • Test-time adaptation (TTA) procedure: After TTA, compilation reduces the adapted model to a two-layer neural network with a skip connection, substantially lowering inference latency.The simplification combines normalization, encoding, and basis weighting into precomputed affine maps.
  • Test-time adaptation (TTA) procedure: Coadaptation provides a non-MAML comparison by jointly optimizing the stimulation context with the remaining parameters during pretraining and adaptation.It uses the same architecture to isolate the effect of MAML.

6) Ablation: pretraining without MAML:

The proposed deployment gradually builds a shared model from early sessions, then uses smaller calibration sets in later sessions; the scheduling policy remains open.

  • Ablation: pretraining without MAML: An example deployment collects larger calibration sets during the first 20 sessions before switching to a meta-learned shared-session TBFM with TTA.The figure contrasts 5k calibration examples initially with 500 after the switch.
  • Ablation: pretraining without MAML: Deployment aims to maximize robustness and minimize calibration-data collection across an experimental campaign.The base model is built during early experiments and used later in the campaign.
  • Ablation: pretraining without MAML: The transition rule depends on the desired accuracy–sample-efficiency trade-off and may require comparing pretrained and single-session models.The paper identifies this choice as an unresolved deployment detail.

7) Deployment:

A deployment strategy uses larger initial calibration sets to pretrain a shared model, then applies smaller calibration sets with test-time adaptation. Pretrained models improve robustness and accuracy on unseen sessions, especially in low-data settings.

  • Deployment strategy: Pretraining after the first N sessions enables smaller calibration sets for later sessions.The deployment procedure uses single-session models initially, then pretrains a shared model for future sessions.
  • Robustness: Pretrained TBFMs significantly improve robustness across all calibration set sizes.Robustness is defined by the left-tail mass of accuracy across held-out sessions and MCCV splits.
  • Robustness: At 1k calibration examples, catastrophic failures fall from 16 of 40 vanilla sessions to 1 of 40 MAML-trained sessions.The catastrophic-failure threshold is test set R2 below 0.05.
  • Uncertainty: Prediction intervals are significantly narrower for MAML-trained models across calibration set sizes.Vanilla models sometimes perform marginally better in individual sessions; at 1k, vanilla outperformed MAML in 8 of 40 sessions.

B. Pretraining and meta-learning significantly increase sample efficiency on unseen sessions

Cross-session pretraining improves sample efficiency on unseen sessions, with MAML outperforming both single-session and co-adapted approaches at moderate calibration sizes. Performance nearly converges by 5k examples, while MAML retains robustness benefits.

  • Sample efficiency: Pretrained models significantly exceed single-session models in sample efficiency, with diminishing returns as calibration size increases.The rates of diminishing returns differ substantially between single-session and pretrained models.
  • Sample efficiency: At 1k calibration examples, MAML achieves mean R2 of 0.397 versus 0.167 for single-session models on held-out sessions.The difference is significant at p = 9.1−5 under the one-sided Wilcoxon signed-rank test.
  • Sample efficiency: At 5k examples, pretrained and single-session performance nearly converges: 0.424 for pretrained models versus 0.398 for single-session models.The difference is statistically insignificant, with p = 0.15.
  • Meta-learning: MAML exceeds co-adaptive pretraining in mean accuracy across all calibration sizes.At 1k, MAML reaches 0.397 R2 versus 0.345 for co-adapted models, with p < 1−5 at all sizes.
  • Meta-learning: MAML also reduces result variance across calibration sizes, although the reduction is not statistically significant.Lower variance indicates more consistent results and may indicate greater robustness.
  • Cross-subject generalization: Pretrained TBFMs successfully generalize between the two subjects in the dataset.Subject G held-out accuracy was effectively identical across models, while Subject J showed a nonsignificant advantage for its subject-specific model.

E. Training time and inference latency slower, but still faster than comparison models

Pretrained TBFMs require more training time and inference computation than vanilla models because of their added complexity, but remain within the stated real-time latency target.

  • Inference latency: Pretrained TBFM GPU inference takes 0.223ms per trial versus 0.121ms for compiled vanilla TBFMs.Despite the slowdown, pretrained inference remains below the 20ms real-time requirement.
  • Training time: MAML pretraining on 20 held-in sessions requires 3.6 hours on average.The authors envision running this training between sessions rather than during a session.
  • Training time: Vanilla model training takes approximately 3 minutes for 5k calibration examples, while TTA takes 13 minutes for an average session with 1k support examples.These timings are reported for the study’s desktop hardware.
  • Training design: MAML is slower here because it initializes only session stimulation embeddings and temporal bases rather than the entire network.This differs from the usual MAML role of initializing a network for rapid few-step adaptation.
  • Ablations: Tanh nonlinearity and weight normalization stabilize training by preventing conflicting gradients from causing parameter-scale oscillations.Without these choices, the basis weight estimator and basis generator fail to stabilize.
  • Ablations: Basis orthonormality regularization lowers held-in accuracy but improves held-out-session accuracy.This suggests the regularizer acts as a constraint supporting generalization.

IV. DISCUSSION

The discussion attributes MAML’s gains to improved adaptation robustness and shared structure across sessions. These benefits reduce calibration demands, but deployment still requires decisions about continued data collection and model refinement.

  • Discussion: Pretrained TBFMs improve sample efficiency and robustness, with MAML adding measurable gains over co-adaptive pretraining.The authors describe this as the first demonstration of practical meta-learning gains for neural stimulation modeling.
  • Robustness gains: At 1k calibration examples, catastrophic failures fall from 16 of 40 vanilla sessions to 1 of 40 MAML-pretrained sessions.MAML prediction intervals are also significantly narrower for calibration sizes up to 2.5k.
  • Mechanisms: MAML optimizes adaptation across randomly sampled support sets, discouraging sensitivity to particular calibration trials.This mechanism is intended to reduce failures under noisy or atypical sessions.
  • Mechanisms: Resting-state autocorrelation provides advance information about session noise before stimulation begins.Conditioning the basis generator on this context allows forecasts to account for session noise level.
  • Practical consequences: MAML-pretrained models achieve comparable or superior accuracy with 50−90% less calibration data.This reduces session time spent on calibration and repeated nontherapeutic stimulation.
  • Cross-session structure: Shared temporal response structure and per-session autoencoders factor common patterns from session-specific spatial differences.The autoencoders handle variable electrode counts and map sessions into a common latent space.

D. MAML vs. Co-Adaptive Pretraining

Cross-session pretraining benefits stimulation forecasting, while MAML further improves adaptation by shaping representations for small calibration sets. The approach remains bounded by dataset scope, runtime, and untested distribution shifts.

  • D. MAML vs. Co-Adaptive Pretraining: MAML-pretrained models outperform co-adapted models across all calibration set sizes, isolating a benefit from the meta-learning objective itself.The comparison uses identical architecture and training corpus, but differs in whether stimulation context parameters are optimized through simulated inner-loop adaptation.
  • D. MAML vs. Co-Adaptive Pretraining: Cross-session pretraining alone substantially improves performance over vanilla single-session training, indicating that previously collected stimulation data is broadly beneficial.The authors attribute the additional MAML improvement to structured inner/outer-loop training that facilitates adaptation rather than merely increasing effective training-set size.
  • D. MAML vs. Co-Adaptive Pretraining: MAML initializes stimulation embeddings so small calibration sets can capture session-specific response properties, while residual-rank constraints keep adaptation well-conditioned.The residual basis matrix uses rank r = 16, limiting degrees of freedom available to the stimulation embedding during test-time adaptation.
  • Limitations: The study’s evidence comes from two rhesus macaques, two cortical regions, and one optogenetic modality, limiting generalizability across species, regions, and stimulation modalities.Electrical stimulation may produce qualitatively different artifact time courses, leaving shared multimodal training unresolved.
  • Limitations: Approximately 13-minute test-time adaptation is acceptable within standard experimental sessions but constrains rapid-iteration paradigms and some clinical applications.Runtime is dominated by autoencoder adaptation in the outer loop; lower-dimensional bottlenecks, PCA initialization, or decoder-only adaptation are proposed as possible reductions.
  • Limitations: Performance under intentional behavioral-state distribution shifts remains unevaluated, including calibration in one state followed by deployment in another.The authors identify multi-state calibration and explicit behavioral-state conditioning as future directions, particularly for adaptive DBS.

H. Clinical translation considerations

The results improve robustness and efficiency relevant to clinical translation, but broader validation and formal safety characterization remain necessary. The model also meets the reported real-time inference target and may suit constrained hardware.

  • H. Clinical translation considerations: Reduced catastrophic failures make MAML-pretrained models less sensitive to calibration-data variability, providing one proxy for robustness to clinical distribution shifts.The authors present this as encouraging but not as full validation across all deployment conditions.
  • H. Clinical translation considerations: A model that performs well on average but fails unpredictably in a minority of sessions is unlikely to satisfy clinical safety and predictability requirements.The observed improvement in tail behavior is described as a step toward, rather than fulfillment of, formal clinical evaluation requirements.
  • H. Clinical translation considerations: 0.223ms GPU inference latency is within the stated 20ms target and faster than some existing clinical systems.The shallow post-adaptation computation and architectural simplicity may support implementation on low-power embedded processors.
  • H. Clinical translation considerations: Inter-session variability across only two animals underscores the challenge of generalizing across patients with different anatomies, disease states, and stimulation histories.The authors frame transfer learning as potentially necessary because large per-patient calibration datasets are often infeasible.
  • H. Clinical translation considerations: The work provides early evidence that cross-session stimulation structure can support pretraining, a necessary but insufficient condition for future neural-stimulation foundation models.The proposed path forward requires heterogeneous datasets spanning modalities, brain regions, and species, plus standardized benchmarks.

VI. APPENDIX

The appendix specifies the TBFM architecture, dataset configuration, and optimization schedules used for meta-learning and test-time adaptation. These settings include shared bases, session embeddings, and explicit inner- and outer-loop controls.

  • VI. APPENDIX: The TBFM configuration uses latent dimension 96, stimulation covariate dimension 3, 100 bases, and a three-layer basis generator.It also specifies latent dimensions of 16 for the basis generator and 32 for the meta projection.
  • VI. APPENDIX: The model uses shared temporal bases, tanh basis weights, a 16-rank meta-learning residual module, and a linear channel autoencoder.The configuration lists a residual MLP hidden dimension of 16 and an autoencoder input dimension of 96.
  • VI. APPENDIX: Each trial contains 184 bins, including a 20-bin runway and a 164-bin forecast horizon, with 5,000 training trials and 2,500 test trials.The protocol holds out 15 sessions and uses an IQR multiplier of 10.0 for outlier handling, with at most five outliers per trial.
  • VI. APPENDIX: Outer-loop training runs for 12,001 steps with batch size 500 per session and gradient clipping at ℓ2 norm 2.0.The basis-weight estimator is updated three times per basis-generator update, using a learning rate of 4×10^-4.
  • VI. APPENDIX: The inner loop uses 20 adaptation steps, a support set of 300, and stimulation-embedding ℓ2 regularization of 5×10^-4.An additional 1,000 tail inner steps refit the stimulation embedding on the full support set after the final outer TTA loop.
  • VI. APPENDIX: TTA runs for 7,001 steps across support sizes of 500, 1,000, 2,500, and 5,000, with basis-generator fine-tuning enabled.The protocol allows adaptation of up to 20 sessions and does not fine-tune basis weights.
Loading 2608.26649v1…