Adaptive Protein Tokenization and Latent Flow Matching for Controllable Protein Design
적응형 단백질 토큰화(Adaptive Protein Tokenization, APT)와 잠재 플로우 매칭(Latent Flow Matching)은 단백질의 복잡한 3차원 구조와 서열 정보를 이산적인 잠재 공간(Discrete Latent Space)으로 압축하고, 이를 속도 벡터장 기반의 생성 모델을 통해 제어 가능하게 설계하는 최첨단 생성 AI 프레임워크입니다. APT는 기존의 국소적 토크나이저와 달리 전역적인 단백질 정보를 담는 토큰을 생성하며, 유한 스칼라 양자화(FSQ)와 플로우 매칭을 결합하여 학습 안정성과 생성 효율성을 극대화합니다.
학습 가능한 코드북 대신 고정된 수학적 격자 구조를 사용하여 코드북 붕괴(Codebook Collapse)를 방지하는 이산화 기술입니다.
노이즈와 데이터 분포 사이의 경로를 따라 속도 벡터장을 학습하여, 확산 모델보다 직관적이고 효율적인 샘플링을 가능케 합니다.
생성 초기에는 강력한 조건부 가이드를, 후기에는 무조건부 생성을 적용하여 단백질 구조의 자연스러움을 보존하는 기법입니다.
3차원 좌표 복원 없이 잠재 토큰 단계에서 보상 함수를 통해 최적의 후보군을 탐색하는 추론 단계 스케일링 기술입니다.
단백질 설계 분야는 거대 언어 모델(LLM)과 확산 모델의 결합으로 급격한 발전을 이루었습니다. 특히 APT 연구(Dilip et al., 2026)는 전역적 특징을 반영하는 토큰을 통해 설계의 유연성과 정확도를 동시에 확보했습니다.
기존 모델들은 VQ-VAE의 코드북 붕괴 문제와 확산 모델의 느린 샘플링 속도로 인해 실제 신약 개발 파이프라인 적용에 한계가 있었습니다. 특히 단백질 구조는 미세한 좌표 오차에도 설계 가능성(Designability)이 손실되기에 정밀한 제어가 필수적입니다.
APT 모델은 유한 스칼라 양자화를 도입하여, levels=(8, 5, 5, 5)와 같은 고정 격자 구조를 사용해 1,000개 이상의 유효 토큰을 안정적으로 확보합니다.
로켓의 조향처럼 초기에는 강한 가이드를, 후기에는 자연스러운 분포를 따르도록 가중치 함수 $w(t)$를 설계합니다.
($t=0$일 때 최대 가이드, $t=1$일 때 무조건부 생성 전환)
Proteína-Complexa (2026)와 유사하게, 토큰 서열만으로 iPAE나 CATH 클래스를 예측하는 경량 분류기를 활용하여 빔 서치의 효율을 극대화합니다.
결합력이 약한 단백질의 접두사(Prefix)를 기반으로 빔 서치를 수행하여 더 강력한 결합을 가진 화합물 설계.
베타 시트(Beta sheet) 비중이 높거나 특정 CATH 분류에 속하는 새로운 단백질 폴딩 구조 생성.
컴퓨터 비전, 생성 모델 및 강화 학습에서 파생된 핵심 기술 요소들의 심층 탐구
기존의 VQ-VAE는 학습 가능한 코드북 $\{e_1, \ldots, e_K\}$를 사용하여 인코더 출력을 가장 가까운 벡터로 치환합니다. 하지만 이는 특정 코드만 활성화되는 코드북 붕괴 문제를 초래하며, 이를 방지하기 위해 EMA나 복잡한 재초기화 기법이 필요합니다.
반면 FSQ는 고정된 격자를 사용합니다. 예를 들어 levels=(8, 5, 5, 5)는 1,000개의 고정된 지점을 의미하며, 인코더는 이 박스 안의 공간을 자연스럽게 활용하도록 학습됩니다.
두 모델 모두 노이즈($t=0$)와 데이터($t=1$) 사이의 경로를 학습하지만, 표현 방식이 다릅니다. 확산 모델은 노이즈나 스코어를 예측하여 역 SDE를 해결하는 반면, 플로우 매칭은 속도 벡터장 $v_\theta$를 학습합니다.
// Flow Matching Loss Function
APT는 학습 시에는 직선 경로를 따르는 플로우 매칭을 사용하지만, 샘플링 시에는 약간의 무작위성을 부여하는 SDE 솔버를 결합하여 샘플의 다양성을 극대화하는 하이브리드 전략을 취합니다.
단백질 구조는 매우 민감하여 강력한 가이드(CFG)를 지속적으로 걸어줄 경우 구조적 정합성이 깨지기 쉽습니다.
"로켓 발사 시 초기 궤도 수정은 강력하게 하되, 착륙 시에는 조종간을 놓고 자연스러운 흐름에 맡기는 것과 같습니다."
APT의 빔 서치는 모델의 로그 확률을 사용하는 대신, 외부 보상 함수 $R(seq)$를 사용합니다. 핵심 효율성 비결은 보상 함수가 3D 구조가 아닌 잠재 토큰(Latent Tokens) 상에서 직접 작동한다는 점입니다.
토큰 서열만으로 베타 시트 비율을 예측하여 높은 점수의 토큰 조합을 선택.
특정 접힘 구조(Fold) 분류에 속할 확률을 보상으로 주어 설계 방향 제어.
결합 오류(iPAE)가 낮을 것으로 예측되는 토큰으로 서열을 반복 개선.
LLM 추론 단계 스케일링(Inference-time Scaling)과의 연계: 이 접근법은 모델의 가중치를 수정하지 않고도 추론 시 더 많은 계산 자원(더 큰 빔 사이즈, 정교한 보상 모델)을 할당하여 성능을 높이는 최근의 추세와 궤를 같이 합니다.
FSQ, 플로우 매칭, 분류기 어닐링, 잠재 빔 서치는 독립적인 기술이 아니라 '전역적이고 적응적인 토큰'이라는 APT의 핵심 철학을 지탱하는 유기적인 시스템입니다. 향후에는 단백질을 넘어 RNA, 대사체 등 다양한 생물학적 데이터를 통합하는 파운데이션 모델로의 진화가 기대됩니다.