Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

BrainDistill: Implantable Motor Decoding with Task-Specific Knowledge Distillation

Yuhan Xie, Jinhan Liu, Xiaoyong Ni, Fei Tan, Icare Sakr, Thibault Collin, Shiqi Sun, Alejandro Rodriguez Guajardo, Demon Fanny, Charles-francois Vincent Latchoumane, Henri Lorach, Jocelyne Bloch, Gregoire Courtine, Mahsa Shoaran

TL;DR

BrainDistill tackles implantable BCI constraints by combining a lightweight neural decoder (IND) with a task-specific knowledge distillation (TSKD) framework that compresses a large teacher’s embeddings into a task-relevant subspace. A novel Task-Specific Ratio (TSR) metric guides projection quality, enabling effective distillation even with limited recalibration data, and a quantization-aware training (QAT) regime supports integer-only inference for energy-efficient, chip-ready deployment. Across diverse ECoG datasets and recording modalities, BrainDistill outperforms traditional decoders and other distillation methods, with particular strength in few-shot calibration and cross-session/generalization. The approach demonstrates on-chip feasibility through 8-bit quantization, modest power costs, and applicability to EEG and spike-based data, signaling a practical path toward clinical, implantable neural decoding systems.

Abstract

Transformer-based neural decoders with large parameter counts, pre-trained on large-scale datasets, have recently outperformed classical machine learning models and small neural networks on brain-computer interface (BCI) tasks. However, their large parameter counts and high computational demands hinder deployment in power-constrained implantable systems. To address this challenge, we introduce BrainDistill, a novel implantable motor decoding pipeline that integrates an implantable neural decoder (IND) with a task-specific knowledge distillation (TSKD) framework. Unlike standard feature distillation methods that attempt to preserve teacher representations in full, TSKD explicitly prioritizes features critical for decoding through supervised projection. Across multiple neural datasets, IND consistently outperforms prior neural decoders on motor decoding tasks, while its TSKD-distilled variant further surpasses alternative distillation methods in few-shot calibration settings. Finally, we present a quantization-aware training scheme that enables integer-only inference with activation clipping ranges learned during training. The quantized IND enables deployment under the strict power constraints of implantable BCIs with minimal performance loss.

BrainDistill: Implantable Motor Decoding with Task-Specific Knowledge Distillation

TL;DR

BrainDistill tackles implantable BCI constraints by combining a lightweight neural decoder (IND) with a task-specific knowledge distillation (TSKD) framework that compresses a large teacher’s embeddings into a task-relevant subspace. A novel Task-Specific Ratio (TSR) metric guides projection quality, enabling effective distillation even with limited recalibration data, and a quantization-aware training (QAT) regime supports integer-only inference for energy-efficient, chip-ready deployment. Across diverse ECoG datasets and recording modalities, BrainDistill outperforms traditional decoders and other distillation methods, with particular strength in few-shot calibration and cross-session/generalization. The approach demonstrates on-chip feasibility through 8-bit quantization, modest power costs, and applicability to EEG and spike-based data, signaling a practical path toward clinical, implantable neural decoding systems.

Abstract

Transformer-based neural decoders with large parameter counts, pre-trained on large-scale datasets, have recently outperformed classical machine learning models and small neural networks on brain-computer interface (BCI) tasks. However, their large parameter counts and high computational demands hinder deployment in power-constrained implantable systems. To address this challenge, we introduce BrainDistill, a novel implantable motor decoding pipeline that integrates an implantable neural decoder (IND) with a task-specific knowledge distillation (TSKD) framework. Unlike standard feature distillation methods that attempt to preserve teacher representations in full, TSKD explicitly prioritizes features critical for decoding through supervised projection. Across multiple neural datasets, IND consistently outperforms prior neural decoders on motor decoding tasks, while its TSKD-distilled variant further surpasses alternative distillation methods in few-shot calibration settings. Finally, we present a quantization-aware training scheme that enables integer-only inference with activation clipping ranges learned during training. The quantized IND enables deployment under the strict power constraints of implantable BCIs with minimal performance loss.
Paper Structure (34 sections, 1 theorem, 18 equations, 7 figures, 18 tables)

This paper contains 34 sections, 1 theorem, 18 equations, 7 figures, 18 tables.

Table of Contents

  1. Introduction
  2. Method
  3. Problem Definition
  4. Task-Specific Knowledge Distillation
  5. Evaluatinig Teacher-Student Projection
  6. IND Architecture
  7. Tokenization
  8. Linear Attention
  9. Quantization-Aware Training with Integer-Only Inference
  10. Experiments
  11. Experimental Settings
  12. Distilled Multi-Class Motor Decoding
  13. Measuring Teacher-Student Projection
  14. Motor Decoding without Teacher Model
  15. Feasibility of On-Chip Implementation
  16. ...and 19 more sections

Key Result

Proposition 2.1

The compression error $\epsilon_{\text{compress}}(U) = \mathbb{E}_{z\sim \mathcal{Z}_{\mathcal{T}}}\|W_{\mathcal{T}}^\top z - (P U)^\top z\|^2$ reaches its minimum $\epsilon_{\text{compress}}(U^*) = \left\|\left(I-\Pi_{\mathcal{U}}^{(\Sigma)}\right) W_{\mathcal{T}}\right\|_{\Sigma}^2$ where $U^* = \

Figures (7)

  • Figure 1: BrainDistill pipeline: (a) Training and testing paradigms of BrainDistill. IND distills from a pre-trained large neural model using small-volume recalibration data via TSKD, and then performs online motor decoding. (b) The two-step TSKD process. First, teacher embeddings are compressed in a supervised manner by a projector; second, the fixed projector is used to align student embeddings with the task space. (c) Task-Specific Projection. TSKD projects teacher embeddings into a task-specific subspace, ensuring that critical task-relevant information is well preserved.
  • Figure 2: Visualization of different projections on Human-C: (a) UMAP visualization of teacher embeddings. (b) UMAP visualization of teacher embeddings projected by $P^*$. (c) UMAP visualization of teacher embeddings projected by PCA. (d) UMAP visualization of teacher embeddings projected by a random orthogonal matrix. Each color denotes a movement, and the density centers of each class are highlighted. (e) The cumulative value $\left\|\Pi_{\mathcal{U}}^{(\Sigma)} W_{\mathcal{T}}\right\|_{\Sigma}^2$ for each projection, computed over 6 task-space dimensions and accumulated. Higher values indicate that the projected features are more likely to yield better decoding performance. For reference, the cumulative $\left\| W_{\mathcal{T}}\right\|_{\Sigma}^2$ of the original feature space is shown in blue, corresponding to TSR = 1.
  • Figure 3: Results of movement regression on Monkey-R. (a) $R^2$ values of IND and different baseline decoders across seven test sessions. (b) $R^2$ values of IND and its variants with different tokenization methods or model architectures; exact $R^2$ values are reported in Appendix \ref{['monkey-r-appendix']}. (c) Illustration of decoding output: comparison of IND and CWT+RNN on a 40-second segment from day 6, with corresponding $R^2$ values.
  • Figure 4: Display of the user interface to visualize online decoding results.
  • Figure 5: Wavelet feature of subject 4 from Human-D. The training set and test set are projected to two two-dimensional space via UMAP respectively.
  • ...and 2 more figures

Theorems & Definitions (1)

  • Proposition 2.1