REVBRAIN: Real-time Probabilistic EEG Forecasting Using Residual Vector Quantization
Authors: Hanna Pankka, Jaakko Lehtinen, Timo Roine
Organizations: Department of Neuroscience and Biomedical Engineering, Aalto University School of Science, Finland · Department of Computer Science, Aalto University School of Science, Finland · NVIDIA, Finland · School of Smart Solutions, Metropolia University of Applied Sciences, Finland
Real-time forecasting of resting-state EEG signals is required by applications like brain-state-dependent stimulation, where computer generated actions are coupled with ongoing neural oscillations. Recently, the deep learning model WaveNet was proposed for the forecasting task. It shows superior accuracy compared to the prevailing method, the linear autoregressive (AR) model, and provides a built-in uncertainty measure due to its probabilistic modeling framework. However, inference using WaveNet is too slow for real-time use due to its sample-by-sample autoregressive operation: a single forecast requires evaluating the model tens or hundreds of times depending on the length of the prediction window. In this paper, we present REVBRAIN, a probabilistic deep learning method that produces forecasts in real time by generating the forecast over the whole prediction horizon at once. This is achieved by tokenizing the output space using residual vector quantization (RVQ) and generating forecasts over the items in the tokenized space. We tested REVBRAIN for forecasting 200 ms of EEG time series in the alpha-frequency band (8-13 Hz) and compared our results against WaveNet model. We found that REVBRAIN achieved real-time inference on a laptop GPU and an over 50× speedup compared to WaveNet, while reaching superior prediction accuracy. By achieving real-time inference speed, REVBRAIN enables incorporating accurate and probabilistic EEG forecasting into BCI applications, opening up new ways for investigating the functioning of the human brain.
Figures & tables
Figure 1: Illustration of producing a quantized representation x^ of an EEG signal x using M=4 codebook levels. The length k of the signals is 100 (equaling to 200 ms at our 500 Hz sampling rate), i.e. the same we use for the prediction horizon. The original and residual signals (in yellow) are mapped to the closest codebook elements (green signals) at the respective codebook levels. First, we find the closest match (highlighted with red) for x from the first codebook ( CB1 ). Then, we calculate the difference between the closest match x^1 and x to form the first residual signal x2 that is then correspondingly mapped to the closest item in CB2 . This process is repeated until we have selected one codebook item from all four codebooks. The final quantized signal x^ is then formed as the sum of the selected items. The rightmost column shows x^ (at the bottom) as well as the intermediate sums.
Figure 2: REVBRAIN is composed of an encoder and M prediction heads. The encoder turns the input into a latent representation z that the heads use to produce the probability distributions over the codebook indices {i1,…,iM} . The first head produces the first index i1 using only z . The following heads then get as inputs both z and all previously predicted indices (indicated with gray arrows). After all the indices are generated, they are used to retrieve the corresponding waveforms {x^1,…,x^M} from the codebooks, and the final prediction is defined as their sum x=∑m=1Mx^m .
Figure 3: REVBRAIN ( 3(a) ) achieved a notable speed-up over WaveNet in inference latency while ( 3(b) ) also reducing the prediction error. ( 3(a) ) The figure shows median latencies over 100 iterations for generating 1, 10, and 100 samples per input. The speed-up was the greatest when generating 100 samples ( 3170× ) and the smallest when generating just 1 sample ( 52× ). ( 3(b) ) The figure shows the mean absolute errors (MAE) on the test data as functions of time. The predictions for both models were formed as the mean of 10 predicted samples. Due to the delay caused by the band-pass filter, the prediction starts 128 ms in the “past” and the true prediction begins at Time=0 (grey vertical line). The mean absolute errors of the true predictions were 2.4±2.4 µV (REVBRAIN) and 2.6±2.7 µV (WaveNet).
Figure 4: Figure ( 4(a) ) depicts the phase bins used for the phase detection: [315°,45°),[45°,135°),[135°,225°),[225°,315°) , roughly corresponding to a peak, falling flank, trough, and rising flank. Figure ( 4(b) ) shows the results of the probabilistic phase detection. The precision rises and recall falls as the threshold is increased—with thresholds of over 90 % prediction was almost never done, but the few that were received close to perfect precision.
Figure 5: The median latencies over 100 iterations for each step in the real-time usage pipeline (top row) and for the whole pipeline (bottom row). The measured signal is first band-pass filtered on the CPU, and the filtered signal is moved to the GPU. Then, the model inference and decision making steps are performed on the GPU. Finally, the binary decision flag is moved back to the CPU to communicate to the system whether to take action or not. Most of the time was consumed by the inference and decision making steps, while the other steps (band-pass filtering and moving to and from GPU) were extremely fast.
Figure 6: ( 6(a) ) The Transformer head (TF-REVBRAIN) increased the inference latencies compared to the MLP-based heads (REVBRAIN), though also this version of the model achieved real-time inference with the smaller sample sizes. ( 6(b) ) The Transformer head resulted in a marginal improvement of 0.01 µV in mean absolute error (MAE) on the test set compared to REVBRAIN.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Head 1
Head 2
Head 3
Head 4
Head 5
FC(256, 512)
Embed(8)
Embed(8)
Embed(8)
Embed(8)
FC(264, 512)
FC(272, 512)
FC(280, 512)
FC(288, 512)
FC(512, 512)
FC(512, 512)
FC(512, 512)
FC(512, 512)
FC(512, 512)
FC(512, 512)
FC(512, 512)
FC(512, 512)
FC(512, 512)
FC(512, 512)
Appendix
Table 1: Layers and layer dimensions of each prediction head of REVBRAIN. The number after an embedding layer signals the embedding dimension. The values for the fully connected (FC) layers correspond to input and output features, respectively. The number of input features for the first FC layer of each head is composed of the dimension of z plus the embedding dimensions of the current and previous levels.
Levels ( M )
Size ( N )
1
2
3
4
5
6
7
8
9
10
128
1.54
0.88
0.55
0.35
0.23
0.15
0.10
0.07
0.05
0.04
256
1.36
0.72
0.41
0.24
0.15
0.09
0.06
0.04
0.03
0.02
train 512
1.19
0.58
0.31
0.17
0.10
0.06
0.03
0.02
0.01
0.01
1024
1.04
0.46
0.22
0.11
0.06
0.03
0.02
0.01
0.01
0.00
2048
0.89
0.35
0.15
0.07
0.03
0.02
0.01
0.01
0.00
0.00
Appendix
Table 4: RVQ model performance with varying codebook sizes ( N ) and numbers of codebook levels ( M ). The reported values are mean absolute quantization errors (µV) over the whole 200 ms window on the train and validation subjects. Bolded values correspond to the errors of our chosen codebook size ( N=512 ) and number of codebook levels ( M=5 ).
Figure 7: Inference time of REVBRAIN with 100 to 10k samples drawn. With a sampling rate of 500 Hz, REVBRAIN can create close to 2k samples per input in real time.
Device
N samples
A100
H100
H200
1
42.0
26.2
25.8
10
70.1
39.3
37.7
100
282.3
142.0
132.0
Appendix
Table 5: Inference latencies of WaveNet with more powerful GPUs: NVIDIA A100, H100 and H200. The reported values are median latencies in milliseconds over 100 iterations. Before measuring, the model was compiled with torch.compile in “max autotune”-mode (v. 2.10.0, Paszke et al., 2019 ) and ran for 50 warm up rounds.
Figure 8: Examples of 10 predicted samples with ( 8(a) ) and without ( 8(b) ) the input-relative encoding of the outputs. The generated samples without the centering step tend to have small offsets relative to the input in the beginning of the prediction, whereas applying the centering ensures that the start of the forecast is aligned with the input time series.
Time series forecasting predicts future values from ordered observations using statistical, neural, and probabilistic models to support decisions under uncertainty.
Electroencephalography (EEG) is a critical, non-invasive method to monitor electrical brain activity. EEGs can span anywhere from a couple seconds to multiple hours, posing a major hurdle for existing deep learning methods due to two major factors: (1) existing EEG models are predominantly built upon the attention mechanism, incurring quadratic scaling as the sequence length increases, and (2) raw EEG signals must be processed in a sliding-window fashion due to fixed-length input requirements, preventing global understanding of the entire signal. To this extent, we propose CaMBRAIN - the first Causal, Mamba-based state space model (SSM) capable of real-time inference of EEG signals, arguing that bidirectional approaches are needlessly expensive given the causal, unidirectional nature of EEG. However, training such a model is non-trivial, as crucial EEG events can be extremely brief - within fractions of a second - yet separated by long intervals spanning minutes. Current EEG methods use self-supervised objectives that optimize for signal reconstruction, but these are not well suited for streaming SSMs; they fail to explicitly train the hidden state to retain the salient long-range context needed for streaming inference. We therefore introduce a multi-stage self-supervised training pipeline specifically tailored to encourage long-range memory retention and strong performance on EEG signals, while preserving the linear-time complexity of state space models. CaMBRAIN achieves state-of-the-art (SOTA) results across 3 different EEG datasets with >10x higher throughput than existing models, enabling the first model capable of long-range, continuous inference of variable-length EEG signals.
Abhilash Durgam, Nyle Siddiqui, Jeffrey A. Chan-Santiago +3
CRCV, University of Central Florida · Department of MAE, University of Central Florida · Department of Neurology, Loma Linda University
Developing foundation models for electroencephalography (EEG) remains challenging due to the signal's low signal-to-noise ratio and complex spectro-temporal non-stationarity. Existing approaches often overlook the hierarchical latent structure inherent in neural dynamics, leading to suboptimal reconstruction of fine-grained information. In this work, we propose BrainRVQ, a general-purpose EEG foundation model pre-trained on a large-scale corpus of clinical EEG data. Unlike standard masked modeling, BrainRVQ features a Dual-Domain Residual Vector Quantization (DD-RVQ) tokenizer that disentangles temporal waveforms and spectral patterns into hierarchical discrete codes. We further introduce a hierarchical autoregressive pre-training objective that learns to reconstruct these codes in a coarse-to-fine manner, utilizing an importance-guided curriculum masking strategy to prioritize information-rich neural events over background noise. Extensive experiments across 8 diverse downstream datasets demonstrate that BrainRVQ consistently outperforms state-of-the-art baselines, validating its effectiveness in learning robust and generalizable neural representations. Our code and model weights are available:https://github.com/keqicmz/BrainRVQ
Mingzhe Cui, Tao Chen, Yang Jiao +4
State Key Laboratory of Industrial Control Technology, Zhejiang University, Hangzhou, China · Department of Electronics, Information and Bioengineering, Politecnico di Milano, Milan, Italy · Shenzhen Institute of Advanced Technology, Chinese Academy of Sciences, Shenzhen, China
Modeling neural population dynamics is crucial for foundational neuroscientific research and various clinical applications. Conventional latent variable methods typically model continuous brain dynamics through discretizing time with recurrent architecture, which necessarily results in compounded cumulative prediction errors and failure of capturing instantaneous, nonlinear characteristics of EEGs. We propose ODEBRAIN, a Neural ODE latent dynamic forecasting framework to overcome these challenges by integrating spatio-temporal-frequency features into spectral graph nodes, followed by a Neural ODE modeling the continuous latent dynamics. Our design ensures that latent representations can capture stochastic variations of complex brain states at any given time point. Extensive experiments verify that ODEBRAIN can improve significantly over existing methods in forecasting EEG dynamics with enhanced robustness and generalization capabilities.
Haohui Jia, Zheng Chen, Lingwei Zhu +6
Information Science and Techinology, Hokkaido University, Japan · SANKEN, The University of Osaka, Japan · Great Bay University, China +1