Edge computing on TPU for brain implant signal analysis

Abstract

The ever-increasing number of recording sites of silicon-based probes imposes a great challenge for detecting and evaluating single-unit activities in an accurate and efficient manner. Currently separate solutions are available for high precision offline evaluation and separate solutions for embedded systems where computational resources are more limited. We propose a deep learning-based spike sorting system, that utilizes both unsupervised and supervised paradigms to learn a general feature embedding space and detect neural activity in raw data as well as predict the feature vectors for sorting. The unsupervised component uses contrastive learning to extract features from individual waveforms, while the supervised component is based on the MobileNetV2 architecture. One of the key advantages of our system is that it can be trained on multiple, diverse datasets simultaneously, resulting in greater generalizability than previous deep learning-based models. We demonstrate that the proposed model does not only reaches the accuracy of current state-of-art offline spike sorting methods but has the unique potential to run on edge Tensor Processing Units (TPUs), specialized chips designed for artificial intelligence and edge computing. We compare our model performance with state of art solutions on paired datasets as well as on hybrid recordings as well. The herein demonstrated system paves the way to the integration of deep learning-based spike sorting algorithms into wearable electronic devices, which will be a crucial element of high-end brain-computer interfaces.

Keywords: Brain–computer interface; Deep learning; Edge device; Electrophysiology; Feature extraction; Spike sorting.