단백질 확산 모델의 테스트 시간 정렬을 위한 스펙트럴 피드백
Spectral Feedback for Test-Time Alignment of Protein Diffusion Models
arXiv확산 모델의 생성을 단방향으로 처리하던 기존 방식과 달리, 스펙트럴 피드백은 반복적인 '피드백 루프'를 통해 생성 오류를 스스로 교정하는 새로운 방법을 제시했습니다.
- 이 알고리즘을 단백질 역접힘에 적용한 결과, 최고 성능 모델 대비 안정적인 단백질 생성률을 최대 32.3%까지 끌어올리는 높은 효율성을 입증했습니다.
- 특정 모델 구조에 구애받지 않는 범용성 덕분에, 단백질 구조 예측 등 고난도 과학 및 공학 분야의 AI 성능 개선에 폭넓게 활용될 잠재력이 큽니다.
- 토큰 간 상호 의존성을 효율적으로 처리하는 것이 핵심이며, 이 기술은 다양한 생성 모델의 정렬(Alignment) 성능을 근본적으로 향상시킬 수 있습니다.
기존의 이산 (discrete diffusion models) 기반 보상 최대화 정렬 방법들은 추론 과정을 단방향으로 처리하여, 원치 않는 선택을 되돌리거나 재검토할 수 있는 메커니즘이 부족했습니다. 본 연구에서 제안하는 스펙트럴 피드백은 '피드백 루프'를 통해 편집 위치(edit-positions)를 선택하고 반복적으로 모델의 생성을 스스로 교정할 수 있게 합니다. 이 접근 방식은 이미지 편집 방법과 유사하게 마스크 구조를 활용하여 토큰을 재마스킹하고 재샘플링하는 방식으로 작동합니다.
이 기술은 단백질 역접힘(protein inverse folding)에 적용되었으며, 편집 효과가 상호 의존적이라는 문제점을 해결하기 위해 edit-set의 값 함수가 희소 푸리에 표현(sparse Fourier representations)을 가진다는 점을 발견했습니다. 이러한 구조 덕분에 스펙트럴 피드백은 편집 위치 선택에 대한 값 함수를 효율적으로 학습하고 최적화할 수 있습니다. 이 알고리즘은 사전 훈련된 모델, 테스트 시간 정렬된 모델, 된 모델 등 어떤 확산 모델에도 적용 가능하며 근본적인 생성 과정을 수정하지 않고도 성능을 개선합니다.
실제 단백질 안정성 보상 오라클(protein stability reward oracle)을 이용한 역접힘에 적용했을 때, 스펙트럴 피드백은 높은 효율성을 입증했습니다. 사전 훈련된 모델의 경우 안정적인 단백질 생성률이 32.3% 증가했으며, Best-of-10 방식에서는 24.8% 증가했고, 최신 RL 미세 조정 확산 모델 대비로는 5.8% 성능 향상을 달성했습니다.
용어 풀이
- 확산 모델
- 잡음에서 시작해 조금씩 다듬으며 이미지나 영상을 만들어 내는 생성 모델.
- 토큰
- AI 모델이 글을 처리하는 단위. 단어보다 작은 조각이며 사용량과 요금을 셀 때 기준이 돼요.
- 미세 조정
- 이미 학습된 모델을 특정 목적의 데이터로 조금 더 학습시키는 일.