EEGAAD
EEG foundation model 이식 및 AAD finetuning 전략 탐색
EEG 기반 Auditory Attention Decoding(AAD) 과제에서 기존 EEG foundation model을 현재 태스크에 붙여보며 finetuning 전략을 조정한 내부 연구 트랙이다. 내가 맡은 작업은 pretrained weight 이식, sampling rate 불일치 처리, task-specific head 설계, selective unfreezing 기반 2-phase finetuning, 그리고 cross-session setting에서의 domain adaptation 방향 정리였다.
문제 설정
과제는 EEG 신호로 청각 주의 방향을 분류하는 문제였다.
- Task 1: cross-subject 학습에 포함되지 않은 피험자에 대해 일반화해야 하는 setting
- Task 2: cross-session audio-only 조건에서 학습하고, audio-visual 조건에서 추론하는 setting
Task 2는 이름만 보면 session 차이처럼 보이지만, 실제로는 단순 session shift보다 도메인 차이가 더 크게 걸리는 문제였다.
내가 본 구간은 baseline을 반복하는 쪽이 아니라, 기존 EEG foundation model을 현재 과제에 어떻게 붙일지, 그리고 어느 범위까지 finetuning하는 것이 맞는지를 확인하는 쪽이었다.
작업 범위
작업은 대부분 비공개 repository의 notebook_ty/ 아래에서 진행했다.
- baseline 코드 구조 확인
- CBraMod pretrained weight 이식
- sampling rate mismatch 처리용 모듈 설계
- task-specific classifier head 설계
- 2-phase finetuning 구현
- unfreezing 비율 비교 실험
- Task 2용 domain adaptation 방향 정리 및 문서화
모델만 교체한 것이 아니라, foundation model을 현재 EEG 태스크 조건에 맞게 붙이기 위한 어댑터와 학습 전략을 같이 만졌다.
CBraMod 이식
backbone으로 사용한 것은 CBraMod였다. AAD 전용 모델은 아니고, 일반적인 EEG representation learning 쪽에서 학습된 foundation model이다. 현재 보고 있는 FM 계열 모델들 중 성능이 강한 편이고, 채널 불균형 처리 측면도 괜찮다고 판단해 기준 backbone으로 잡았다.
가장 먼저 걸린 부분은 sampling rate였다.
- 대회 데이터: 128 Hz
- CBraMod pretraining 조건: 200 Hz
이 차이를 무시하지 않고, 중간에 TemporalResampler를 두는 방식으로 처리했다. 1D 선형 보간으로 128 Hz를 200 Hz로 맞췄고, 입력 shape도 (B, C, T)와 (B, C, 1, T) 둘 다 받을 수 있게 구현했다.
다만 이 부분은 업샘플링이기 때문에 한계가 분명하다. 원래 없던 정보를 복원하는 건 아니고, pretrained 입력 조건에 맞추기 위한 타협에 가깝다. 다음에 다시 진행한다면 128 Hz 조건으로 pretrain을 다시 돌리거나, 가능하면 원본 데이터를 맞춰서 pretraining 조건 자체를 다시 가져가는 쪽을 먼저 볼 것 같다. 현재로서는 128 Hz 기준 재학습이 가장 현실적이다.
task-specific head
backbone 뒤에는 현재 과제에 맞는 classifier head를 따로 붙였다.
AdaptiveAvgPool2dLinearELUDropoutLinear(2)
AAD는 최종적으로 2-class 분류 문제이기 때문에, 기존 출력부를 그대로 쓰기보다 현재 문제에 맞는 head를 다시 다는 편이 맞다고 봤다. backbone의 proj_out도 nn.Identity()로 바꿔서 feature map을 그대로 흘려보내도록 처리했다.
finetuning 전략
학습은 처음부터 전체 layer를 다 여는 방식으로 가지 않았다. pretrained feature가 빠르게 무너질 가능성이 있다고 봤기 때문이다.
그래서 2-phase finetuning으로 갔다.
phase 1
- classifier만 먼저 학습
- cold start 형태
- 초반 5 epoch 사용
phase 2
- encoder 상위 일부 layer만 선택적으로 unfreeze
- backbone learning rate는 classifier보다 낮게 설정
- pretrained feature를 최대한 보존한 채 현재 태스크에 적응시키는 방향
언프리즈 비율은 고정하지 않고 10% / 30% / 50%로 나눠 비교했다. 감으로 정하지 않고, 실제로 어느 정도 열었을 때 더 낫게 붙는지 확인하는 쪽으로 봤다.
최종적으로는 상위 50%를 여는 쪽을 더 낫게 봤고, 최종 실험도 그 방향으로 정리했다. 공개 범위에서는 구체 수치 대신 “baseline 대비 성능 개선을 확인했다” 정도로 두는 것이 맞다.
같이 넣은 학습 요소
Task 1 최종 학습 로직에는 아래 요소들을 같이 넣었다.
- Mixup
- Contrastive loss
- selective unfreezing
- differential learning rate
backbone만 바꾼 것이 아니라, 현재 setting에서 pretrained representation을 덜 무너뜨리면서 일반화를 확보하려는 방향으로 학습 구조도 같이 조정했다.
Task 2: 문제를 다시 본 부분
Task 2에서는 모델 구조를 바꾸는 것보다, 문제를 먼저 다시 정의하는 쪽이 더 중요했다.
초기 접근에서는 audio-only와 audio-visual 데이터를 혼합해서 다루는 방식의 오류가 있었는데, 이 상태에서는 domain shift를 제대로 해석하기 어렵다고 판단했다. 그래서 먼저 아래처럼 나눠서 봤다.
- source domain: audio-only
- target domain: audio-visual
같은 EEG 분류 문제로 두기보다, 명시적인 domain adaptation 문제로 다시 본 것이다.
그 위에서 적용 방향으로 정리한 것은 아래였다.
- CORAL loss
- Contrastive loss
- Mixup
- entropy minimization
여기서는 기법 자체보다, source와 target을 먼저 분리해서 해석하는 쪽이 더 중요했다.
내가 맡은 부분
이 프로젝트에서 맡은 역할은 아래와 같다.
- 기존 EEG foundation model(CBraMod) 이식
- sampling rate mismatch 해결용 TemporalResampler 설계
- task-specific classifier head 설계
- cold start → selective unfreezing 구조의 2-phase finetuning 구현
- unfreezing 비율 비교 실험(10/30/50)
- Task 2의 domain shift 구조 재정의
- domain adaptation 기법 적용 방향 설계 및 문서화
처음부터 새 모델을 만드는 프로젝트라기보다, 기존 representation을 현재 태스크에 맞게 옮기고, 어느 조건에서 실제로 잘 붙는지를 확인하는 작업에 가까웠다.
정리
이 프로젝트에서 직접 걸렸던 부분은 결국 아래였다.
- EEG foundation model은 다른 EEG 태스크에도 바로 쓸 수 있는가
- 그 과정에서 sampling rate, head 구조, unfreezing 범위가 얼마나 민감하게 작동하는가
- Task 2처럼 겉보기에는 같은 분류 문제여도, 실제로는 domain shift가 더 본질적인 문제가 아닌가
특히 Task 2를 보면서는, 단순히 구조를 바꾸는 것보다 더 일반화된 EEG representation 자체가 필요하다는 쪽으로 생각이 갔다. 더 다양한 EEG 조건에서 학습된 foundation model이 있어야 이런 setting에서 덜 흔들릴 가능성이 크다고 봤다.