CABT: Causal-Aware Bidirectional Transformer BCI Decoder

CABT is a high-performance Causal-Aware Bidirectional Transformer neural decoder designed to translate intracortical brain signals into text during handwriting movements. Developed for multi-electrode array recordings, CABT serves as a modern, attention-based replacement for legacy Recurrent Neural Network (GRU) decoders, achieving superior Character Error Rates (CER) and faster execution speeds.

This repository contains clean, modular, and reproducible scripts optimized to train and evaluate the CABT decoder on Apple Silicon (M-series) Macs and distributed Linux clusters using Slurm.

---

Architectural Features

Unlike standard sequential models, CABT processes brain signal dynamics in parallel using a specialized transformer architecture tailored for neural interfaces:

1. Temporal Patch Embedding High-frequency neural recordings contain substantial correlation across consecutive frames. CABT groups adjacent time steps (e.g., 5 frames / 100ms) into unified patch tokens. This downsamples the sequence length by $5\times$, saving attention memory computation, and captures localized temporal trends before attention modeling.

2. Causal-Aware Bidirectional Attention Mask Standard Transformers are either fully causal (only attending to the past) or fully bidirectional (attending to the entire sequence). In online BCI decoding, a physical delay (lag) is introduced between neural activity and target characters. CABT implements a lookahead attention mask that allows each time step $i$ to attend to all past steps and a configurable window of future steps $L$ matching the decoder delay: $$M_{i, j} = \begin{cases} 0 & \text{if } j \le i + L \ -\infty & \text{if } j > i + L \end{cases}$$ This respects causality limits in physical deployment while maximizing prediction accuracy by leveraging lookahead context.

3. Rotary Position Encoding (RoPE) Instead of adding absolute positional embeddings to inputs, CABT embeds relative position information by rotating query and key representations in the complex plane. This improves generalization across varied sentence lengths.

4. GeGLU Feed-Forward Blocks CABT replaces the standard feed-forward block with a Gated Linear Unit utilizing GELU activation (GeGLU), improving model capacity and routing stability.

---

Repository Structure

• cabt.py: Core decoder modules (Temporal Embedding, MHA with RoPE, GeGLU, Transformer block) and training/inference execution wrappers.
• data_loader.py: Pre-processing pipeline which handles block-based baselining, time-series binning (defaulting to 20ms bins), padding, and session partitions.
• evaluation.py: Decodes raw logits into text sentences based on rising-edge threshold crossings of the start signal, and computes CER and WER via Levenshtein edit distance.
• vocab.py: Vocabulary mappings, character lists (31 tokens), Kaldi lexicons, and duration values.
• run_train.py: Training orchestration script that maps sessions, configures parameters, and manages checkpoints.
• run_infer.py: Runs evaluations on test sets, computes global CER/WER metrics, and saves probability output matrices.
• slurm_job.sh: Batch cluster worker submission template.

---

Setup & Installation

Prerequisite Environment

CABT is built in pure TensorFlow 2.12.0 and has been optimized for both CUDA-enabled Linux servers and Apple Silicon Macs (Metal Acceleration) using the bci_tf2 Conda environment.

To set up the environment on macOS:

conda create -n bci_tf2 python=3.10
conda activate bci_tf2
conda install -c apple tensorflow-deps
pip install tensorflow-macos tensorflow-metal
pip install scipy numpy pickle-mixin

---

Training and Inference

1. Training CABT

Modify the directory configurations in run_train.py to point to your local dataset directories, then execute:

python run_train.py

Note: Training checks for existing checkpoint indexes in the output directory. If a crash or pause occurs, running the script again will automatically resume from the last saved step.

2. Evaluation

To run inference over the validation blocks and print comparative targets vs predictions:

python run_infer.py

3. Submitting Slurm Jobs

To launch training in the background on a GPU cluster partition:

sbatch slurm_job.sh

---

Local Compilation & Integration Test

A verification script is included to test model initialization, Gaussian smoothing, patching, masking, RoPE, GeGLU, loss computation, backpropagation, and checkpoint writing on synthetic data tensors:

python verify_cabt.py

The local compile checks function completely without requiring external Matlab dataset files.

---

References & Acknowledgments

This project is built upon and inspired by the dataset, research, and baseline configurations from:

• Willett, F. R., Avansino, D. T., Cunningham, J. P., Shenoy, K. V., & Hochberg, L. R. (2021). High-performance brain-to-text communication via handwriting. Nature, 593(7858), 249-254.
• We extend our sincere gratitude to the authors for making their dataset public. While the data preprocessing, HMM labels alignment, and language model evaluation pipelines in this repo align with the methodology established in their paper, the recurrent neural network baseline (GRU) has been replaced by our custom-designed CABT (Causal-Aware Bidirectional Transformer) decoder architecture.