Flow Matching for Few-Trial Neural Adaptation with Stable Latent Dynamics

We sincerely thank you for your careful review and recognition of our work. Below, we provide a point-by-point reply to your concerns. Due to the limit, all figures prefixed with 'R' below are available in the external link [https://drive.google.com/file/d/129vv370SF4RLanLj92-lzh_vMmkeDCve/view?pli=1].

Claims And Evidence:

• About why the stability of the flow is made out to be so critical for transfer and a more thorough ablation

Thanks for raising the unclear point. The stability of flow models enhances zero-shot performance, making it effective for few-trial adaptation. Specifically, this stability regulates feature deviations, preserving essential semantic information across domains. As illustrated in Fig.R4, it ensures that latent factors with similar labels flow toward consistent semantic representations, even under input drifts in target signals. Empirical validation through zero-shot transfer performance (original Table S8 on Page 18) further demonstrates the benefits of this stability.

Additionally, we conducted a more thorough ablation study to further validate this effect. Since stability is maintained by activation functions and scale coefficients, we ablated these two components (FDA-al and FDA-sc, respectively) to violate the assumptions of stability. As shown in Fig.R5, the distribution of maximum Lyapunov exponents (MLE) for FDA-al and FDA-sc indicates that both variants frequently exhibit positive MLEs, signifying instability. Furthermore, the corresponding results for zero-shot and few-trial performance, evaluated using MMD-based alignment on the CO-M and RT-M datasets, are summarized in the table below. Consistent with their reduced stability, FDA-al and FDA-sc demonstrated substantially degraded transfer performance, which is comparable to the baselines. This result suggests that instability is the factor contributing to their poor transfer ability.

Comparison of average $R^2$ scores (%) across sessions for FDA-al, FDA-sc, and FDA-MMD on the CO-M, and RT-M datasets ($r$ = 0, 0.02).

Other Comments Or Suggestions:

• I don't think that presenting the results in a table, as in Table 1, is the most effective way of communicating them in this case. While this is not uncommon in ML papers, I find this table to be too busy to parse. I'd suggest something more like a box or violin plot may help. Maybe with just the most important results, and then the rest could be moved to the supplementary material.

Thanks a lot for the valuable suggestion. We will improve the presentation of our results by using clearer violin plots and relocating some results to the supplementary material in our revision.

Questions For Authors:

• About the easy way to ablate the stability property of FDA and show that it transfers worse in some way when those assumptions are violated

Yes, we found that FDA exhibited significantly worse zero-shot and few-trial transfer performance when these assumptions were violated. Specifically, since stability is governed by activation functions and scale coefficients, we ablated these two components (FDA-al and FDA-sc, respectively) to affect the stability property of FDA. More detailed results are provided above in the above response.

We sincerely hope that these responses may address your concerns. We believe that our novel FDA framework will be of significant interest to the ICML community, given its potential impact on few-trial neural alignment and real-world BCI reliability. Could you please consider raising the scores? We look forward to your further feedback. Thank you in advance.

We sincerely thank you for your careful review and recognition of our work. Below, we provide a point-by-point reply to your concerns.

• Additional comparison with NDT-2

We conducted an additional comparison against NDT-2 by pre-training it on a single session with supervised readout training and evaluating its zero-shot performance without alignment. The average $R^2$ scores on the CO-C and CO-M datasets are presented in the table below. While NDT-2 achieved high decoding performance in the pre-trained session, its zero-shot performance degraded significantly. This decline may be attributed to its reliance on larger datasets to achieve robust zero-shot transfer.

Average $R^2$ scores (%) of NDT-2

• About the supervised alignment by NDT-2 and MtM

Thank you for the valuable comment. We agree that alignment of large-scale models isn't necessary in the cases we have presented in this study.

What we meant is that the finetuning of MtM and NDT-2 can begin with self-supervised techniques, such as masked reconstruction. Then, for different downstream tasks, specific decoders are trained using only a few target labels, while all other weights remain fixed. This supervised alignment enables rapid transfer to tasks that were not encountered during pre-training.

• About other BCI tasks and cross-task/subject adaptation

Thanks for the comment. We agree that this presents interesting directions for further study.

Questions:

• About the drop of $R^2$ with more finetuning samples

The curve in Fig.4(b) and Fig.S6 is based on a single random run, and the observed decrease may be due to the small trial gap (~2 trials between adjacent points). Moreover, the overall trend shows an increase in $R^2$ as the number of finetuning samples grows, as demonstrated in Fig.3(c) on Page 6.

• About the choice of window size w

Thanks. We conducted a grid search to determine the appropriate context window size $w$. As shown in the Table S11 on Page 21, balancing performance and computational efficiency, we selected 5 as the default size for the CO-M and RT-M.

• About better $R^2$ on RT than CO tasks

The better performance on RT tasks may be attributed to differences in signal quality between subjects. For the same subject (Monkey M), performance on CO tasks is better than on RT tasks, as shown in the original Fig.3(b) on Page 6.

• About the desirable level of diversity in trials

We conducted additional analyses on selected sessions of CO-M. The diversity was categorized into three levels: small (1–2 distinct targets), medium (3–4 targets), and large (all distinct targets). The average $R^2$ for FDA-MMD is presented in the tables below. Our results indicate that a medium level of diversity achieves desirable performance, while too little diversity negatively impacts performance.

Average $R^2$ scores (%) for FDA-MMD on CO-M

• About the linear decoder

The linear decoder was mentioned in Line 162 (right) on Page 3, and will be made clearer in the revision.

• About details of hardware and reported computational efficiency

The hardware used was NVIDIA GeForce RTX 3080 Ti (12GB). The training time reported in Table S9 on Page 20 summarizes the time per epoch for both pre-training and fine-tuning. A more detailed comparison of total time(s) is presented below.

• About the efficiency of inference

We conducted further analysis of FDA's inference time on an NVIDIA GeForce GTX 1080 Ti (11GB). As presented in the table below, average inference time for a single window is approximately 4 ms, making it suitable for real-time applications.

Inference time of FDA (ms)

We hope that these responses may address your concerns. We believe that our novel FDA framework will be of significant interest to the ICML community, given its potential impact on few-trial neural alignment and real-world BCI reliability. Could you please consider raising the scores? We look forward to your further feedback. Thank you in advance.

We sincerely appreciate your thorough review and recognition of our work. Below we provide a detailed response to your concerns, with reference figures in the supplementary material [https://drive.google.com/file/d/129vv370SF4RLanLj92-lzh_vMmkeDCve/view?pli=1], indicated by Fig.R.

Weaknesses

• Model interpretability

The latent features effectively capture the evolved state within neural dynamics. We employ a transformer-based network $f_\alpha$ to extract latent features from signal windows. Given the challenges in interpreting the meaning of these features for realistic scenarios, our current focus is on the effectiveness of FDA for few-trial adaptation. We demonstrate the successful extraction of the Lorenz attractor from synthetic spiking data using our method below. We will explore the relationship between the latent features and underlying neural dynamics across various cases in future.

• Synthetic data

We conducted experiments to evaluate the recovery of ground-truth latent variables in synthetic data. Following method in (Kapoor, Jaivardhan, et al. Latent diffusion for neural spiking data. NeurIPS 2024: 118119), we used the Lorenz attractor as the latent dynamics. We simulated firing rates as an affine transformation of the 3D latent variables into a 96-dimensional space, then sampled spike trains from a Poisson distribution.

As shown in the table, our FDA successfully recovered the latent dynamics from synthetic spiking data. The visualizations of our decoded 3D trajectories in Fig.R1 confirmed that FDA effectively captured the neural dynamics.

Average $R^2$ (%) on the recovered latent variables

• About z (0-1)

The evolution of z only represents the iterative learning process. The temporal evolution of neural dynamics is characterized by the shift of short-term windows(as explained in Line 106-108 on Page 2).

Comments:

• We will clarify the multi-dimensional labels in the revision.

• Here $\tau$ represents the iterative learning steps instead of the temporal evolution of neural signals. To avoid confusion, we therefore used a distinct notation $\tau$ .

• The scale coefficient represents the shift parameter of z generated by the MLP for predicting the velocity field.

• Yes, it is all $l_2$ distance and will be standardized.

• We find that our flow model can be well-suited to both few-trial adaptation in BCIs and non-BCI recordings. The flow process can include feedback as conditional features to guide the flow of subsequent windows.

Questions:

• Session definition

The sessions used in experiments were from the same subject. The recorded neurons across sessions largely originate from the same subset. Moreover, as shown in Fig.R2, the transformer-based $f_{\alpha}$ is adaptable to varying numbers of tokens, corresponding to different channel counts.

• Non-overlapping subset

Given the widespread application of transformers for universal representation from neurons across regions, we believe our model has the potential to be extended to cases with non-overlapping neuron subsets across sessions. We agree that this is an interesting direction for future study.

• One step later

Yes, "one step later" refers to the next sample based on the sampling rate. The overlap between windows depends on the sampling rate, ensuring consistent time steps between the decoded and ground-truth variables.

• $\tau$ (0-1)

As explained above, $\tau$ represents the iterative learning step of z to obatin latent features, differing from the temporal evolution of neural signals.

• Choice of $\eta$

$\eta$ is pre-initialized using Xavier initialization. We tested various choices via differnt random seeds, and observed no significant impact on the final results (as shown in Table 1 on Page 7).

• Simultaneous neurons

The current model directly utilized temporal structures from single-channel recordings as input tokens for conditional features. Therefore, it is preferable for neurons to be recorded simultaneously within each session.

• Poisson statistics

As shown in the table above, the model demonstrated adaptability to Poisson synthetic data with various firing rates. Furthermore, as illustrated by the coefficient of variation (CV) distribution in Fig.R3, some real data we used exhibits characteristics of Poisson distributions.

• Monkey C

The figure for Monkey C is shown in Fig.S1(a).

We sincerely hope that these responses may address your concerns. We believe that our novel FDA framework will be of significant interest to the ICML community, given its potential impact on few-trial neural alignment and real-world BCI reliability. Could you please consider raising the scores? We look forward to your further feedback. Thank you in advance.

We sincerely thank you for your careful review and recognition of our work. Below, we provide a point-by-point reply to your concerns. Due to the limit, all figures prefixed with 'R' below are available in the external link[https://drive.google.com/file/d/129vv370SF4RLanLj92-lzh_vMmkeDCve/view?pli=1].

Methods And Evaluation Criteria

• About the motivation

Our FDA successfully addresses the issues of significant degradation in zero-shot performance and instability in few-trial alignment, achieving stable latent feature extraction and reliable alignment with few trials.

On one hand, flow-based learning regulates feature deviations through stable latent dynamics, preserving essential semantic information from the source domain. This mechanism enforces bounded variations in latent features under input target perturbations, preserving semantic information in neural representations. Empirical validation of zero-shot transfer performance (original Table S8 on Page 18) highlights the advantage of this stability, further enhancing efficient few-shot adaptation.

On the other hand, by leveraging the unique properties of flows, FDA guarantees stable few-trial fine-tuning. The use of explicit log-probability objectives and flexible latent-space modeling enables stable parameter gradients, effectively preventing the catastrophic overfitting typically seen in few-trial adaptation.

Experimental Designs Or Analyses

• About more intuitive visualization on the flow matching process

Thank you for the suggestion. We have provided a more intuitive visualization of the flow matching process in Fig.R4. As shown in this figure, our flow matching process learns latent variables from neural signals in a coarse-to-fine manner, differing from conventional one-step extraction.

The flow transforms noisy variables $\mathbf{z}(0)$ into neural representations $\mathbf{z}(1)$, guided by conditional features derived from neural signals. Simultaneously, the corresponding prior distribution $p_0$ is transformed into the target distribution $p_1$. The stable latent dynamics of the learning process ensure that latent factors with similar labels flow toward consistent neural representations, even when guided by shifted neural signals.

Essential References Not Discussed

• About the missing related works that specifically tackling the variability of brain signal decoding

Thanks a lot for the suggestion. We will include more related works that specifically tackle the variability of brain signal decoding in the revision.

Other Comments Or Suggestions:

• About more annotations in Fig. 2. For example, what does X^S and X^T respectively stands for. Similarly for C^S and C^T. And what is the relationship between C^S and v_{\theta}?

Thank you for the valuable comment. We have added more annotations, and the revised version is provided as Fig.R2. Here, $\mathbf{x}^S$ and $\mathbf{x}^T$ represent input tokens derived from short-term windows of the source and target domains, with each token corresponding to spikes from a single channel. Additionally, $\mathbf{c}^S$ and $\mathbf{c}^T$ denote the output conditional features of the transformer-based $f_{\alpha}$, which are utilized to guide the flow of $\mathbf{z}$.

As illustrated in Fig.R4, we parameterize the guided velocity field of $\mathbf{z}$ as $v_{\theta}$, which determines the velocity and direction of $\mathbf{z}$. $\mathbf{c}^S$ further serves as the input to the neural network responsible for predicting $v_{\theta}$.

• About more explanations on why target distribution p1 representing the desired neural representation could be defined by a random variable z^S(1)?

Thank you for raising this point. To achieve optimal decoding performance, the target distribution is preferable to be defined by a random variable, which can then be transformed into the ground-truth labels through the decoders. When weights of the linear decoder are fixed, the desired $\mathbf{z}^S(1)$ can be obtained via a linear transformation using the inverse of the decoder's weight matrix. The neural representation thus defined can subsequently be transformed into the desired labels through the decoder. Related explanation will be included in the revision.

We sincerely hope that these responses may address your concerns. We believe that our novel FDA framework will be of significant interest to the ICML community, given its potential impact on few-trial neural alignment and real-world BCI reliability. Could you please consider raising the scores? We look forward to your further feedback. Thank you in advance.