[논문 리뷰] AffinityNet: semi-supervised few-shot learning for disease type prediction

Tianle Ma, Aidong Zhang|arXiv (Cornell University)|2018. 05. 22.

Machine Learning in Bioinformatics참고 문헌 24인용 수 8

한 줄 요약

AffinityNet는 고차원적이고 표본 수가 적은 유전체 데이터로부터 강력한 표현을 학습하기 위해 스택형 k-최근접 이웃(kNN) 어텐션 풀링 레이어를 사용하는 반감독적 소수의 시도 학습 모델로, 암 유형 예측에 활용된다. 이 모델은 전통적인 딥러닝 모델과 베이스라인 코ックス 모델을 능가하며, 훈련 데이터의 40%만을 사용해 신장암 데이터에서 0.729의 일치지수를 기록하여 낮은 데이터 환경에서 뛰어난 일반화 성능을 입증한다.

ABSTRACT

While deep learning has achieved great success in computer vision and many other fields, currently it does not work very well on patient genomic data with the "big p, small N" problem (i.e., a relatively small number of samples with high-dimensional features). In order to make deep learning work with a small amount of training data, we have to design new models that facilitate few-shot learning. Here we present the Affinity Network Model (AffinityNet), a data efficient deep learning model that can learn from a limited number of training examples and generalize well. The backbone of the AffinityNet model consists of stacked k-Nearest-Neighbor (kNN) attention pooling layers. The kNN attention pooling layer is a generalization of the Graph Attention Model (GAM), and can be applied to not only graphs but also any set of objects regardless of whether a graph is given or not. As a new deep learning module, kNN attention pooling layers can be plugged into any neural network model just like convolutional layers. As a simple special case of kNN attention pooling layer, feature attention layer can directly select important features that are useful for classification tasks. Experiments on both synthetic data and cancer genomic data from TCGA projects show that our AffinityNet model has better generalization power than conventional neural network models with little training data. The code is freely available at https://github.com/BeautyOfWeb/AffinityNet .

연구 동기 및 목표

• 고차원적 특징과 적은 표본 수로 인해 딥러닝 성능이 저하되는 유전체 데이터에서의 '큰 p, 작은 N' 문제를 해결하기 위해.
• 생물의학 연구에서 흔한 소수의 시도 학습 환경에서 잘 일반화되는 데이터 효율적인 딥러닝 모델을 개발하기 위해.
• 사전에 정의된 그래프 구조가 필요 없이 환자와 같은 복잡하고 이질적인 대상에 대해 반감독적 표현 학습을 가능하게 하기 위해.
• 입력 특징에 대한 중요도 가중치를 부여하는 학습 가능한 특징 어텐션 레이어를 통해 딥러닝에 자동 특징 선택 기능을 통합하기 위해.
• 제한된 레이블된 데이터와 풍부한 비레이블된 데이터를 활용하여 암 유전체학에서 생존 예측 및 질병 아형 분류를 향상시키기 위해.

제안 방법

• 핵심 구성 요소는 kNN 어텐션 풀링 레이어로, 어텐션 기반 기법을 사용해 동적으로 유사도 그래프를 구성하고 각 객체의 k개의 최근접 이웃으로부터 표현을 풀링함으로써 군집화 유사 정규화 효과를 제공한다.
• kNN 어텐션 풀링 레이어는 그래프가 없는 데이터에 대해서도 적용 가능한 그래프 어텐션 모델(GAM)의 일반화된 형태로, 그래프의 존재 여부와 관계없이 어떤 객체 집합에도 적용 가능하다.
• 다중 kNN 어텐션 풀링 레이어를 스택하여 깊이 있는 계층적 표현을 학습하며, CNN의 합성곱 레이어와 유사한 플러그인 모듈 기능을 수행한다.
• 특징 어텐션 레이어는 kNN 어テン션 풀링의 특수한 경우로, 역전파를 통해 정규화되고 비음수인 가중치를 입력 특징에 대해 학습함으로써 엔드 투 엔드 특징 선택을 가능하게 한다.
• 모델는 반감독 방식으로 훈련되어 레이블된 데이터와 비레이블된 데이터를 모두 활용해 표현 학습과 일반화 성능을 향상시킨다.
• 프레임워크는 PyTorch로 구현되었으며, 재현성과 확장성을 위해 GitHub에 공개되어 있다.

실험 결과

연구 질문

• RQ1 고차원적이고 표본 수가 적은 유전체 데이터에서 소수의 레이블된 예시만으로도 딥러닝 모델이 효과적으로 학습할 수 있는가?
• RQ2 사전에 정의된 그래프 구조가 없을 경우 어텐션 기반 기법을 어떻게 변형하여 표현을 학습할 수 있는가?
• RQ3 기존 신경망과 비교해 kNN 기반 어텐션 풀링이 질병 아형 분류의 소수의 시도 학습에서 일반화 성능을 향상시킬 수 있는가?
• RQ4 모델이 비레이블된 데이터를 얼마나 효과적으로 활용하여 생존 예측 작업의 성능을 향상시킬 수 있는가?
• RQ5 학습 가능한 특징 어텐션 레이어의 통합이 유전체 분류에서 모델 성능 향상과 해석 가능성 향상에 기여하는가?

주요 결과

• AffinityNet은 훈련 데이터의 40%만을 사용해 신장암 데이터에서 평균 일치지수 0.729를 기록했으며, 베이스라인 코ックス 모델(0.618)을 크게 능가했다.
• 훈련 데이터를 10%로 제한했을 때도 AffinityNet은 일치지수 0.694를 기록하여 낮은 데이터 환경에서도 뛰어난 일반화 성능을 입증했다.
• 모델이 예측한 위험군(저위험, 중간위험, 고위험) 간의 로그랭 테스트 p-값은 6.7×10⁻¹⁶이었으며, 생존 결과를 분리하는 데서 기존의 알려진 질병 유형(p=1.5×10⁻⁵)보다 뛰어난 성능을 보였다.
• 합성 및 TCGA 암 유전체 데이터셋에서 AffinityNet은 제한된 훈련 데이터로도 기존 신경망 모델보다 더 우수한 일반화 성능을 보였다.
• 특징 어텐션 레이어는 생물학적으로 관련성이 있는 특징을 성공적으로 식별하고 강조하여 학습 중에 자동이고 미분 가능한 특징 선택을 가능하게 했다.
• kNN 어텐션 풀링 레이어는 유사한 객체들 간에 유사한 표현을 강제함으로써 소수의 데이터 환경에서 과적합을 억제하는 암묵적 정규화 기능을 수행했다.