Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

RL-BioAug: Label-Efficient Reinforcement Learning for Self-Supervised EEG Representation Learning

Cheol-Hui Lee, Hwa-Yeon Lee, Dong-Joo Kim

TL;DR

The paper tackles the data-augmentation bottleneck in EEG contrastive learning caused by non-stationarity and labeling costs. It introduces RL-BioAug, a label-efficient reinforcement learning framework with a Transformer-based policy that autonomously selects strong augmentations to optimize self-supervised EEG representations, using only $10\%$ labeled data for policy guidance. Rewards are derived from a Soft-KNN consistency score on a reference embedding space, and policy updates follow REINFORCE++ with advantage and entropy regularization, ensuring stable, exploration-conscious learning. Experiments on Sleep-EDFX and CHB-MIT demonstrate substantial improvements over random augmentation and reveal task-specific augmentation strategies, highlighting the approach's potential for autonomous, label-efficient EEG analysis.

Abstract

The quality of data augmentation serves as a critical determinant for the performance of contrastive learning in EEG tasks. Although this paradigm is promising for utilizing unlabeled data, static or random augmentation strategies often fail to preserve intrinsic information due to the non-stationarity of EEG signals where statistical properties change over time. To address this, we propose RL-BioAug, a framework that leverages a label-efficient reinforcement learning (RL) agent to autonomously determine optimal augmentation policies. While utilizing only a minimal fraction (10\%) of labeled data to guide the agent's policy, our method enables the encoder to learn robust representations in a strictly self-supervised manner. Experimental results demonstrate that RL-BioAug significantly outperforms the random selection strategy, achieving substantial improvements of 9.69\% and 8.80\% in Macro-F1 score on the Sleep-EDFX and CHB-MIT datasets, respectively. Notably, this agent mainly chose optimal strategies for each task -- for example, Time Masking with a 62\% probability for sleep stage classification and Crop \& Resize with a 77\% probability for seizure detection. Our framework suggests its potential to replace conventional heuristic-based augmentations and establish a new autonomous paradigm for data augmentation. The source code is available at \href{https://github.com/dlcjfgmlnasa/RL-BioAug}{https://github.com/dlcjfgmlnasa/RL-BioAug}.

RL-BioAug: Label-Efficient Reinforcement Learning for Self-Supervised EEG Representation Learning

TL;DR

The paper tackles the data-augmentation bottleneck in EEG contrastive learning caused by non-stationarity and labeling costs. It introduces RL-BioAug, a label-efficient reinforcement learning framework with a Transformer-based policy that autonomously selects strong augmentations to optimize self-supervised EEG representations, using only labeled data for policy guidance. Rewards are derived from a Soft-KNN consistency score on a reference embedding space, and policy updates follow REINFORCE++ with advantage and entropy regularization, ensuring stable, exploration-conscious learning. Experiments on Sleep-EDFX and CHB-MIT demonstrate substantial improvements over random augmentation and reveal task-specific augmentation strategies, highlighting the approach's potential for autonomous, label-efficient EEG analysis.

Abstract

The quality of data augmentation serves as a critical determinant for the performance of contrastive learning in EEG tasks. Although this paradigm is promising for utilizing unlabeled data, static or random augmentation strategies often fail to preserve intrinsic information due to the non-stationarity of EEG signals where statistical properties change over time. To address this, we propose RL-BioAug, a framework that leverages a label-efficient reinforcement learning (RL) agent to autonomously determine optimal augmentation policies. While utilizing only a minimal fraction (10\%) of labeled data to guide the agent's policy, our method enables the encoder to learn robust representations in a strictly self-supervised manner. Experimental results demonstrate that RL-BioAug significantly outperforms the random selection strategy, achieving substantial improvements of 9.69\% and 8.80\% in Macro-F1 score on the Sleep-EDFX and CHB-MIT datasets, respectively. Notably, this agent mainly chose optimal strategies for each task -- for example, Time Masking with a 62\% probability for sleep stage classification and Crop \& Resize with a 77\% probability for seizure detection. Our framework suggests its potential to replace conventional heuristic-based augmentations and establish a new autonomous paradigm for data augmentation. The source code is available at \href{https://github.com/dlcjfgmlnasa/RL-BioAug}{https://github.com/dlcjfgmlnasa/RL-BioAug}.
Paper Structure (27 sections, 8 equations, 3 figures, 3 tables)

This paper contains 27 sections, 8 equations, 3 figures, 3 tables.

Table of Contents

  1. INTRODUCTION
  2. METHODS
  3. Cooperative Optimization Loop
  4. Non-causal Environments and Context-based Decision
  5. Reinforcement Learning Problem Formulation
  6. Agent: Transformer-based Policy Network
  7. Self-Supervised Learning Step: Encoder Update via Contrastive Loss
  8. Reinforcement Learning Step: Agent Update via Policy Gradient
  9. Reward Definition via Soft-KNN
  10. REINFORCE++ Optimization
  11. Variance Reduction via Advantage
  12. Entropy Regularization
  13. Final Objective Function
  14. EXPERIMENTS
  15. Dataset Description
  16. ...and 12 more sections

Figures (3)

  • Figure 1: Overview of the RL-BioAug Framework. The process consists of four distinct phases: (1) State Observation, where the state is extracted via a frozen encoder; (2) Action, where the agent selects an optimal strong augmentation; (3) SSL Step, which updates the encoder using contrastive loss between weak and strong views; and (4) RL Step, which evaluates the representation quality using the Soft-KNN consistency score on a reference set to reward the agent.
  • Figure 2: Experiment Protocol.
  • Figure 3: Training dynamics of the data augmentation policy on (A) the Sleep-EDFX dataset and (B) the CHB-MIT dataset.