EEG 신호를 활용한 청소년 ADHD 진단

요약

Github: ADHD-EEG-ViT

주의력결핍 과잉행동장애(Attention deficit / hyperactivity disorder, ADHD)는 아동 및 청소년기에 흔히 발생하는 신경발달장애예요. 조기에 적절히 진단하고 개입하지 않으면 학업 성취, 사회적 관계, 정서 발달에 장기적으로 부정적 영향을 줄 수 있죠. 여러 뇌파 신호 중 EEG(Electroencephalogram)는 비침습적으로 뇌 활동을 측정할 수 있어 신경과학 연구와 임상 진단 분야에서 널리 활용되고 있어요.

본 연구에서는 EEG 신호를 활용해 청소년 ADHD 진단을 돕는 딥러닝 모델을 설계하고 실험했어요. Vision-Transformer(A. Dosovitskiy et al., 2021)와 EEG-Transformer(Y. He et al., 2023)의 아이디어를 바탕으로 transformer 기반 모델을 구현했어요. IEEE에서 제공하는 “EEG Data ADHD-Control Children” 데이터셋을 활용하여 0.972의 높은 정확도를 달성했답니다.

본 모델의 주요 장점은 다음과 같아요.

• 별도의 복잡한 전처리 과정 없이 end-to-end 학습이 가능해요.
• Mixed precision 기법을 활용해 학습 속도를 높임과 동시에 높은 정확도를 기록했어요.
• Embedding layer를 조정하여 다른 EEG 데이터셋에도 쉽게 적용할 수 있는 확장성을 고려했어요.

다만, 학습 과정에서 모델의 과적합(overfitting) 현상을 발견했어요. 이는 제한된 데이터셋으로 인한 것으로 보이며, 향후 추가 데이터 확보나 데이터 증강(data augmentation) 기법을 통해 모델의 안정성(robustness)을 개선할 수 있을 것으로 기대돼요.

선행 연구 요약

EEG-Transformer(Y. He et al., 2023)는 Transformer(A. Vaswani et al., 2017)의 Self-Attention 구조를 차용한 EEG 분석 모델을 제안했어요. 특히 Attention blocks, Residual connection, Normalization이 데이터 분석에 핵심적인 역할을 한다고 강조했죠.

또 다른 접근 방식으로는 CNN을 활용한 연구들이 활발히 진행되었어요. 그 중 비교적 최근 연구(M. Y. Esas and F. Latifoğlu)에서는 Robust Local Mode Decomposition (RLMD), Variational Model Decomposition (VMD)와 같은 전처리 기법을 사용해 데이터의 특징을 추출하고, 이를 CNN에 통과시켜 ADHD 여부를 판단해요.

데이터셋

본 연구에서는 IEEE data-port의 “EEG Data ADHD-Control Children” 데이터셋을 활용했어요.

데이터 주요 특징은 다음과 같아요.

• 총 121명의 참가자로, ADHD 그룹 61명과 건강한 대조군 60명으로 구성되어 있어요.
• 7-12세 어린이를 대상으로 진행된 연구로, 전문가가 DSM-IV 기준에 따라 ADHD를 진단했어요.
• EEG 신호는 10-20 standard*에 따라 19개 채널(Fz, Cz, Pz, C3, T3, C4, T4, Fp1, Fp2, F3, F4, F7, F8, P3, P4, T5, T6, O1, O2)로 기록되었으며, 128 Hz의 sampling frequency를 사용했어요.
• 실험 방법은 아이들에게 캐릭터 사진을 보여주고 캐릭터의 수를 세도록 하는 과제를 제시했어요.

데이터에 대한 정보가 명확하고 두 그룹 간 균형이 잘 맞춰져 있어 연구에 적합한 데이터셋이에요.

10-20 standard는 EEG 두피 전극 부착 위치에 대한 국제 표준이에요. 데이터셋의 구체적인 전극 위치는 preprocess.ipynb에서 확인할 수 있어요.

전처리

IEEE 데이터셋은 아이들이 캐릭터 수를 세는 과제 중 측정된 EEG 신호를 포함하고 있어요. 참가자마다 과제 완료 시간이 다르기 때문에 EEG 신호의 길이도 다양해요. 구체적으로 신호 길이는 7,983부터 43,252까지 다양하며, 대체로 ADHD 그룹의 과제 완료 시간이 더 길었어요.

선행 연구(M. Y. Esas and F. Latifoğlu)에서는 데이터 길이를 9,250으로 고정하는 접근 방식을 제안했어요. 이 방법은 여러 장점이 있답니다.

• 대부분의 데이터를 활용할 수 있어 데이터 손실을 최소화해요.
• 충분한 길이를 통해 신호의 맥락 정보를 더 잘 보존할 수 있어요.

일부 연구(D. Tanko et al., 2022)에서는 데이터를 더 짧은 단위로 나누기도 했지만, 본 연구는 Transformer의 강점을 활용하기 위해 맥락 정보를 충분히 담을 수 있는 방식으로 데이터를 샘플링하고자 해요. 따라서 다음과 같은 방식으로 데이터를 처리했어요.

• 9,250보다 짧은 데이터는 분석에서 제외해요.
• 9,250보다 긴 데이터는 9,250 단위로 나눠요. (예: 19,000 길이의 데이터는 9,250 * 2 + 500으로 나누어 2개의 subset만 사용하고 나머지 500 길이는 버려요.)

전처리 과정을 거친 데이터는 학습 데이터와 테스트 데이터로 분리했어요. 전체 데이터의 80%를 학습에, 20%를 테스트에 할당하여 총 138개의 학습 데이터와 36개의 테스트 데이터를 확보했답니다.

모델 설계

본 모델은 크게 Embedding과 Transformer 두 부분으로 구성되어 있어요. Embedding은 Vision Transformer(A. Dosovitskiy et al., 2021)를 참고했으며, Transformer는 EEG-Transformer(Y. He et al., 2023)의 구조를 기반으로 일부 파라미터를 수정했어요.

Convolutional Embedding

EEG-Transformer는 데이터 채널의 차원을 Transformer 입력 차원으로 사용하기 때문에 다른 데이터에 대해 적용이 어려워요. 특히 본 연구와 같이 채널의 크기가 작을 때 성능에 치명적인 영향을 줄 수 있죠. (원본 모델은 56 채널 데이터를 사용했어요.)

Vision Transformer 연구에서는 이미지 embedding에 convolution을 활용할 때 공간 정보를 효과적으로 포착할 수 있음을 보여주었어요. 신호 처리 분야에서 CNN이 우수한 성능을 보이고 있다는 점을 고려해, Convolution을 이용한 Embedding layer로 데이터의 차원을 변환하는 접근 방식을 채택했어요.

Embedding된 벡터($z$)는 positional encoding($E_{pos}$)과 합산돼요. Positional encoding은 학습 가능한 파라미터로, 신호의 시간적 정보를 포함하기 위해 도입됐어요.

\[E_x = Conv1d(x)\] \[z = E_x + E_{pos}\]

기존 Vision Transformer와의 주요 차이점은 [CLS] 토큰을 사용하지 않는다는 것이에요. 대신 계산된 모든 feature vector를 분류에 활용하며, 이는 EEG-Transformer의 접근 방식을 최대한 반영한 결과랍니다.

Attention Blocks

Embedding 벡터는 다음과 같은 구조의 Attention block으로 처리돼요.

• Multi-head Attention
• Residual connection + Layer Normalization
• Linear transformation
• Dropout
• Linear transformation
• Residual connection + Layer Normalization

이 과정을 반복하여 EEG 신호의 특징을 추출해요.

Residual connection은 Attention Block의 입력($z$)을 직접적으로 더하는 과정을 의미해요.

\[x_{attn} = Attention(z)\] \[x' = LayerNorm(z + x_{attn})\]

이는 모델이 원본 데이터의 특성을 잘 반영하도록 하기 위해 ResNet에서 제안한 방법이에요. 원본 데이터를 더하는 identity mapping을 통해 모델 가중치를 크게 변형하지 않아도 데이터 특성을 잘 파악하도록 도와줘요.

Classifier

추출된 feature vector는 Global Max Pooling을 통해 차원을 축소해요. 이후 Feed-forward network에 입력되어 최종적으로 ADHD 여부를 분류해요.

구체적인 구현은 ViTransformer에서 확인할 수 있어요.

학습

Colab 환경에서 T4 GPU로 학습을 진행했으며, 구체적인 학습 설정은 다음과 같아요:

• Batch size: 8
• Gradient accumulation: 4 steps
• Cross-entropy loss
• Adam optimizer
• Learning rate: 0.001
• Linear warmup: 30 steps
• Early stopping: 30 step patience
• 5-fold cross validation
• Automatic mixed precision (FP16)

상세한 학습 과정은 ieee_transformer.ipynb에 기록되어 있어요.

결과 분석

Accuracy	Recall	F1-score
0.972	0.952	0.976

모델은 0.972의 높은 정확도를 달성했으며, 약 30 epoch 근처에서 수렴하는 양상을 보였어요.

모델 깊이

Attention block의 차원을 64-128-64로 설정하고 attention head는 4개로 구성했어요. 그리고 이러한 block을 총 4번 반복했죠. 더 깊은 모델 구조는 오히려 성능 저하를 야기했는데, 이는 제한된 데이터셋 규모로 인해 많은 파라미터가 완전히 학습되지 못했기 때문으로 추측돼요.

Mixed Precision

Pytorch의 Auto mixed precision을 활용해 FP32와 FP16 정밀도를 혼합했어요. 이 접근 방식으로 학습 속도를 약 3배 개선했으며 최종 모델 성능에는 영향을 미치지 않았어요.

한계

학습 과정에서 validation loss를 통해 모델의 과적합(overfitting) 현상을 관찰했어요.

Dropout이나 weight decay와 같은 정규화 기법을 적용했음에도 불구하고 과적합 문제를 완전히 해결하지 못했어요. 이는 제한된 데이터셋의 근본적인 한계로 보입니다.

결론

본 연구는 Vision Transformer와 EEG-Transformer의 아이디어를 결합하여 ADHD 진단을 위한 딥러닝 모델을 제시했어요. EEG 신호 분석에 Transformer 아키텍처를 적용함으로써 0.972의 뛰어난 정확도를 달성했답니다.

주요 의의는 다음과 같아요.

1. End-to-end 학습: 별도의 복잡한 전처리 과정 없이 신호의 특징을 추출했으며 추가 데이터셋을 학습하여 모델 성능을 지속적으로 개선할 수 있어요.
2. 높은 성능: 0.972의 정확도와 0.976의 F1 점수를 기록하며 ADHD 진단의 가능성을 보여주었어요.
3. 확장성: Embedding layer 설계를 통해 다양한 EEG 데이터셋에 적용 가능한 모델 구조를 개발했어요.

그러나 연구의 한계 또한 분명합니다. 제한된 데이터셋으로 인한 과적합 문제는 향후 해결해야 할 중요한 과제예요. 추가 데이터 확보, 데이터 증강 기법, 보다 정교한 정규화 방법 등을 통해 모델의 안정성과 일반화 성능을 개선할 수 있을 것입니다.

이 연구는 EEG를 활용한 ADHD 진단 기법을 제시함으로써 수치 데이터를 활용한 객관적인 진단의 가능성을 보여주었어요. 앞으로 더 많은 데이터를 통해 청소년 ADHD 조기 진단에 기여할 수 있을 것으로 기대돼요.

참고 자료

• Y. He et al., “Classification of attention deficit/hyperactivity disorder based on EEG signals using a EEG-Transformer model,” J. Neural Eng., vol. 20, no. 5, Sep. 2023.
• M. Y. Esas and F. Latifoğlu, “Detection of ADHD from EEG signals using new hybrid decomposition and deep learning techniques,” J. Neural Eng., vol. 20, no. 3, Jun. 2023.
• D. Tanko et al., “EPSPatNet86: eight-pointed star pattern learning network for detection ADHD disorder using EEG signals,” Physiol. Meas., vol. 43, no. 3, Apr. 2022.
• A. Dosovitskiy et al., “An image is worth 16×16 words: Transformers for image recognition at scale,” arXiv preprint arXiv:2010.11929, 2021.
• K. He et al., “Deep residual learning for image recognition,” arXiv preprint arXiv:1512.03385, 2015.
• P. Micikevicius et al., “Mixed Precision Training”, arXiv preprint arXiv:1710.03740, 2018.