TRACE: Contrastive learning for multi-trial time series data in neuroscience

We thank the reviewer for their positive assessment of our work and the interesting questions.

(1) Subset-averaging strategy: We acknowledge the reviewers concern about neurons with highly variable or contradictory responses across trials potentially creating misleading positive pairs. TRACE operates under the assumptions that differences between trials are (a) undesirable noise but (b) characteristic for the local noise structure between different neurons. In some cases, between-trial variation may be of interest, e.g., when investigating adaptation over multiple trials. This violates our assumption (a) and TRACE is not the right tool for this type of question. Conversely, if between-trial variation is not of interest and some neuron responds very differently in different trials, this may simply mean that responses are very noisy for this and similar neurons. The resulting diverse positive samples are not misleading, but rather correctly teach TRACE to strongly denoise this part of the data. In fact, TRACE's data trial-driven positive pairs may be less likely to "wash out" meaningful difference than other methods, which operate on averages across all trials and thus do not incorporate any information about the between-trial variation. Standard approaches like t-SNE and UMAP compute Euclidean distances after averaging across all trials, which can incorrectly bias them by uninformative parts of the time series (Sec 5.1). CEED applies generic data augmentations that may or may not be aligned to the true local noise structure of the data. We think this point raised by the reviewer is important and we will incorporate this discussion into Sec 6 of our manuscript for the revision.

(2) Generation of positive samples: We thank the reviewer for pointing out the insufficient description of how we generate positive samples. To clarify, during training, each mini-batch contains subset means from multiple neurons. For each neuron, we create one positive pair consisting of two time series. To create one part of the pair, a random subset of trials of that neuron is selected and averaged. To create the second part of the pair, another non-overlapping subset is averaged. One such pair of subsets is dynamically resampled for each neuron per epoch. We will describe this positive sample generation in a revised Sec 3.2 and the code will also be made publicly available.

We hope this resolves your concerns. Please let us know if you have any concerns left.

(1) Additional neural data set: Based on the reviewer's comment, we applied TRACE to an additional neural dataset with spike responses: the Allen Institute neuropixels 'visual coding' dataset (de Vries et al. 2020; Siegel et al. 2021). It's the largest publicly available dataset of spiking responses across visual brain areas, containing 10,322 neurons (after quality filtering, excluding non-visual neurons) from awake and behaving mice responding to light flashes (450 ms, n_trials=75) and drifting gratings (n_trials=14, 250 ms, 8 orientations, 5 frequencies). We applied TRACE and other methods to binned spike trains at a resolution of 25 ms. The spike train data consists of discrete neural events that follow Poisson-like statistics, where the timing and occurrence of spikes carry important information. TRACE is especially beneficial for such spiking data because preserves the natural statistical structure through trial subset averaging, whereas standard augmentations used by contrastive methods (Gaussian noise, temporal jitter) add continuous noise that destroys the discrete nature of spike trains. We evealuated the performance using standard clustering metrics (ARI, kNN accuracy, Spearman distance correlation) with brain region as class labels and biologically meaningful variables including orientation selective index (OSI), preferred orientation (PO), grating modulation ratio (F1/F0), natural image selectivity (NIS; based responses to natural images not used to create the 2D embedding), behavioral modulation (correlation of firing rate with running speed of the mouse; not used to create the 2D embedding). Alignement of the embedding with biological variables was measures with kNN regression ($R^2$). For PO and behavior, we computed the radial correlation, as it better captured their global structure in the embedding. We also compute the average rank each method had per metric. While standard methods (t-SNE, UMAP) found some structure especially in terms of distance correlation and for OSI, TRACE outperformed them in most metrics and CEED in all of them. The performance of TS2Vec variants was mixed, with good preservation of some structure (OSI, NIS) but very poor global distance correlation. Indeed, the global structure of the TRACE+TS2Vec was highly inconsistent across runs. Overall, TRACE achieved the best rank by far, indicating best overall performance on this dataset. Particularly noteworthy is TRACE's ability to identify image selectivity and modulation by behavior without this information being included during training. Here we report the mean across three seeds:

We believe this additional validation on a large-scale dataset addresses the reviewer's concern about generalization and demonstrates TRACE's broad applicability across neuroscience dataset types.

(2) Novelty of the approach: While we acknowledge that some individual components of our method existed before, we respectfully argue that our core contribution lies in the novel idea of creating positive pairs by subset means over trials. Moreover, TRACE is the first method to combine Cauchy kernel and self-supervision to address specific challenges in neuroscience data visualization. This is not a common combination, as most contrastive methods use cosine similarity, not suitable for visualization. This leads to the following advantages of TRACE over previous methods: (1) Noise robustness: Results on simulated data show that standard methods like t-SNE and UMAP are sensitive to noisy baseline activity, while current contrastive approaches like CEED use augmentations that mask subtle biological differences. Only TRACE was able to successfully separate the cell types by combining trial subset averaging with contrastive learning for direct 2D visualizations (Sec. 5.1, Fig 2). (2) Artifact detection: We showed that TRACE directly identifies experimental artifacts that would would otherwise remain undetected. The artifact island was only clearly visible in the TRACE embedding (Sec 5.3, Fig. 4). Together, these novel contributions make TRACE a valuable tool for the neuroscience community. Current methods either lack the domain-specific design needed for multi-trial data or fail to provide biological interpretability. This gap is addressed through the TRACE framework.

(3) Role of subset means: We take subset means to create diverse positive pairs that are still similar to the inference-time inputs, not primarily to increase SNR (see end of Sec 3.2). Indeed, it is important that the positive pairs are still noisy and do not have the highest possible SNR, because it is precisely their noise structure to which TRACE learns to become invariant. That is the reason why TRACE also works decently when positive pairs are individual trials (Sec 5.2, Fig S2, Table S2). If we instead computed the average across all trials to maximize the SNR, both members of the positive pair would be identical and not provide any learning signal.

(4) Generalizability across domains: While we acknowledge that we present TRACE as a framework particularly suitable for neuroscience data, many time-series datasets have an inherent multi-trial structure or can be reshaped into this format. Often, time-series data contains natural repetitions or can be segmented into comparison units. To give a couple of examples: In sport analytics inertial measurements are often taken during repeated exercises, such as basketball free-throw drills (Hoelzemann et al. 2023, García-de-Villa et al. 2022). The Google Speech Command Dataset (Warden et al. 2018) contains audio recordings of speech commands with 5 trials per speech command and speaker. For medical data such as ectrogardiography one could use different daily cycles or stress tests as trials. For financial market data, trading periods can be used as trials to detect unusual patterns or market anomalies. Another application domain could be climate data where years, seasons, or tidal cycles could be used as trials. We are happy to add this discussion our Sec 6 and hope that this convinces the reviewer that time series data sets from many domains could benefits from TRACE’s ability to identify outliers and subtle pattern differences that standard methods miss.

References: [1] Hoelzemann, A., Romero, J. L., Bock, M., Laerhoven, K. V., and Lv, Q. (2023). Hang-time HAR: A benchmark dataset for basketball activity recognition using wrist-worn inertial sensors. Sensors, 23(13), 5879. [2] García-de-Villa, S., Jiménez-Martín, A., and García-Domínguez, J. J. (2022). A database of physical therapy exercises with variability of execution collected by wearable sensors. Scientific Data, 9(1), 266. [3] Warden, P. (2018). Speech commands: A dataset for limited-vocabulary speech recognition. arXiv preprint arXiv:1804.03209.

Hopefully, we could address your concerns and would appreciate you increasing your score!

(1) Additional neural data set: Based on the reviewer's comment, we applied TRACE to an additional neural dataset with spike responses: the Allen Institute neuropixels 'visual coding' dataset (de Vries et al. 2020; Siegel et al. 2021). It's the largest publicly available dataset of spiking responses across visual brain areas, containing 10,322 neurons (after quality filtering, excluding non-visual neurons) from awake and behaving mice responding to light flashes (450 ms, n_trials=75) and drifting gratings (n_trials=14, 250 ms, 8 orientations, 5 frequencies). We applied TRACE and other methods to binned spike trains at a resolution of 25 ms. The spike train data consists of discrete neural events that follow Poisson-like statistics, where the timing and occurrence of spikes carry important information. TRACE is especially beneficial for such spiking data because preserves the natural statistical structure through trial subset averaging, whereas standard augmentations used by contrastive methods (Gaussian noise, temporal jitter) add continuous noise that destroys the discrete nature of spike trains. We evealuated the performance using standard clustering metrics (ARI, kNN accuracy, Spearman distance correlation) with brain region as class labels and biologically meaningful variables including orientation selective index (OSI), preferred orientation (PO), grating modulation ratio (F1/F0), natural image selectivity (NIS; based responses to natural images not used to create the 2D embedding), behavioral modulation (correlation of firing rate with running speed of the mouse; not used to create the 2D embedding). Alignement of the embedding with biological variables was measures with kNN regression ($R^2$). For PO and behavior, we computed the radial correlation, as it better captured their global structure in the embedding. We also compute the average rank each method had per metric. While standard methods (t-SNE, UMAP) found some structure especially in terms of distance correlation and for OSI, TRACE outperformed them in most metrics and CEED in all of them. The performance of TS2Vec variants was mixed, with good preservation of some structure (OSI, NIS) but very poor global distance correlation. Indeed, the global structure of the TRACE+TS2Vec was highly inconsistent across runs. Overall, TRACE achieved the best rank by far, indicating best overall performance on this dataset. Particularly noteworthy is TRACE's ability to identify image selectivity and modulation by behavior without this information being included during training. Here we report the mean across three seeds:

We believe this additional validation on a large-scale dataset addresses the reviewer's concern about generalization and demonstrates TRACE's broad applicability across neuroscience dataset types.

(2) Dimensionality of the latent space: TRACE is not inherently limited to 2D and the architecture supports any dimensionality. We focus on 2D because our paper is about visualization which requires 2D output independent of the intrinsic dimensionality of the data. Based on the reviewer's comment, we generated representations for 16 and 128 dimensions across three seeds:

kNN accuracy:

Spearman correlation:

TRACE outperforms competitor models in higher dimensions in kNN accuracy and places rank 2 for the Spearman correlation. Notably, TRACE maintains consistent performance across all dimensions, while TS2Vec shows significant degradation in higher dimensions for Spearman correlation, demonstrating TRACE's robustness across different dimensionalities.

(3) Quantitative performance metrics: We appreciate the reviewer's interest in quantitative evaluation. However, we respectfully disagree that our work lacks quantitative performance metrics. We dedicate Section 4.2 specifically to quantitative measures, employing three standard metrics widely used in dimensionality reduction literature: Adjusted Rand Index (ARI), kNN accuracy, and Spearman's rank correlation. For the toy dataset, we additionally provide discriminability in the low-dimensional space (Sec. 5.1, Fig 2). Further, there is an important limitation for our biological dataset: Standard machine learning quantitative metrics assume the existence of ground truth labels, which are rarely available in experimental neuroscience. The class labels were derived from clustering the neural responses themselves (using other methods), making standard clustering evaluation metrics potentially circular. To address this challenge, we additionally used independent biologically meaningful variables such as the ON-OFF index, response-transience index, and recording depth. Recording depth, in particular, represents an experimental variable not derived from neural responses and thus serves as external validation. The strong correlations we observe between our embeddings and these biologically meaningful measures provide an appropriate assessment of method quality for neuroscience applications. We believe our current quantitative evaluation - combining standard dimensionality reduction metrics with biological validation - provides a comprehensive assessment suitable for this applications.

(4) Generalizability across domains: While we acknowledge that we present TRACE as a framework particularly suitable for neuroscience data, many time-series datasets have an inherent multi-trial structure or can be reshaped into this format. Often, time-series data contains natural repetitions or can be segmented into comparison units. To give a couple of examples: In sport analytics inertial measurements are often taken during repeated exercises, such as basketball free-throw drills (Hoelzemann et al. 2023, García-de-Villa et al. 2022). The Google Speech Command Dataset (Warden et al. 2018) contains audio recordings of speech commands with 5 trials per speech command and speaker. For medical data such as ectrogardiography one could use different daily cycles or stress tests as trials. For financial market data, trading periods can be used as trials to detect unusual patterns or market anomalies. Another application domain could be climate data where years, seasons, or tidal cycles could be used as trials. We are happy to add this discussion our Sec. 6 and hope that this convinces the reviewer that time-series data sets from many domains could benefits from TRACE’s ability to identify outliers and subtle pattern differences that standard methods miss. Additionally, you asked whether TRACE works with other loss functions. Our setup requires a contrastive loss function and we chose the most prominent one, InfoNCE. We plan to explore other options, such as Noise Contrastive Estimation, Negative Sampling, or Triplet loss in future work.

References: [1] de Vries, Saskia EJ, et al. "A large-scale standardized physiological survey reveals functional organization of the mouse visual cortex." Nature neuroscience 23.1 (2020): 138-151. [2] Siegle, Joshua H., et al. "Survey of spiking in the mouse visual system reveals functional hierarchy." Nature 592.7852 (2021): 86-92. [3] Hoelzemann, Alexander, et al. "Hang-time HAR: A benchmark dataset for basketball activity recognition using wrist-worn inertial sensors." Sensors 23.13 (2023): 5879. [4] García-de-Villa, Sara, Ana Jiménez-Martín, and Juan Jesús García-Domínguez. "A database of physical therapy exercises with variability of execution collected by wearable sensors." Scientific Data 9.1 (2022): 266. [5] Warden, Pete. "Speech commands: A dataset for limited-vocabulary speech recognition." arXiv preprint arXiv:1804.03209 (2018).

Hopefully, we could address your concerns and would appreciate you increasing your score!