Multi-Objective Protein Design via Memory-Aware Test-Time Scaling in Diffusion Models

Motivation

• Known: diffusion model 기반 단백질 설계는 자연 단백질의 복잡한 분포를 잘 포착하며, 최근에는 retraining 비용을 줄이기 위해 reward function을 활용한 test-time scaling(예: classifier guidance, SMC-based guidance, value-based sampling, iterative refinement)이 단일 목표 최적화에 효과적으로 사용되고 있다.
• Gap: 기존 test-time diffusion 방법들은 (1) 상호작용 히스토리로부터 학습하지 못해 반복적인 설계 오류를 유발하고, (2) 성공 사례(high-reward sample)에만 의존해 실패 사례의 결정 경계 정보를 무시하며, (3) 다중 목표를 단순 선형 가중합으로 처리해 dominant-objective bias를 일으키는 등 multi-objective trade-off를 균형있게 다루지 못하는 한계가 있다.
• Why: 합성생물학 및 신약개발에서 실제 단백질 응용은 안정성, 소수성, 표면 접근성, 면역원성, 용해도 등 상충되는 다중 속성을 동시에 만족해야 하므로, retraining 없이 유연하게 multi-objective 균형을 맞출 수 있는 test-time 방법이 실용적으로 매우 중요하다.
• Approach: MOMST는 memory bank를 통해 과거 반복 히스토리에서 일반화 가능한 추론 경험을 추출하고, self-contrastive learning으로 성공/실패 사례 모두로부터 residue-level 선호 신호를 도출하며, inference-time Pareto alignment 전략으로 다중 목표 간 충돌을 해결하는 방식으로 diffusion model의 test-time 생성을 안내한다.

Achievement

1. 메모리 인지형 test-time scaling 패러다임 제안: 생성 과정을 insight-driven evolutionary process로 재구성하는 MOMST 프레임워크를 제안함.
2. Memory Bank와 Self-contrastive Learning 결합: elite buffer(고보상 히스토리)와 sliding-window exploration buffer(최신 탐색)를 이용한 dual-memory 설계로 반복 오류를 방지하고, 성공/실패 사례를 대조하여 residue-level 선호 신호를 도출함.
3. Inference-time Pareto alignment 전략 제시: 가중합 기반 local sampling, Pareto Frontier를 통한 dominated solution 제거, Pareto sampling을 통한 global resampling을 결합하여 다중 목표 간 균형을 확보함.
4. 단일/다중 목표 과제 모두에서 우수한 성능 입증: 다양한 실험을 통해 MOMST가 기존 방법 대비 뛰어난 성능을 보임을 확인함.

How

• Memory Bank: elite buffer(과거 고보상 샘플 저장)와 sliding-window exploration buffer(최근 탐색 이력 저장)의 dual-memory 구조로 구성되어 모델의 진행 상황에 맞춘 가이던스를 제공
• Self-contrastive Learning: 고보상 시퀀스와 실패 시도(low-reward)를 대조하여 residue 수준의 상대적 선호 신호를 도출, 불안정성 또는 기능 손실과 연관된 residue를 억제
• Inference-time Pareto Alignment: (i) 가중합 기반 local sampling으로 구조적 다양성 보존, (ii) Pareto Frontier로 dominated solution 제거 및 최적 trade-off 유지, (iii) Pareto sampling으로 global resampling을 수행해 모든 목표의 lower bound를 최대화
• 전체 과정은 closed-loop cycle로 구성되어 반복적으로 memory 갱신과 preference 신호 도출, Pareto 정렬을 수행함 (Figure 1)

Originality

• 기존 test-time diffusion 방법들이 memoryless하게 매 스텝의 즉각적 reward만 활용하는 것과 달리, 과거 상호작용 히스토리 전체를 활용하는 memory bank 개념을 도입함
• 성공 사례에만 의존하던 기존 reward 신호 설계와 달리, 실패 사례까지 활용하는 self-contrastive learning을 통해 residue-level 상대적 선호 신호를 도출하는 새로운 방식을 제시함
• 다중 목표를 단순 선형 가중합으로 처리하던 기존 방법의 한계를 극복하기 위해, Pareto Frontier와 Pareto sampling을 결합한 inference-time Pareto alignment 전략을 새롭게 제안함
• 에이전트의 memory-aware test-time scaling 아이디어를 단백질 설계 diffusion model에 최초로 접목함

Limitation & Further Study

• 논문 발췌본에서는 memory bank의 크기, elite/exploration buffer의 구체적 갱신 규칙, Pareto alignment의 계산 비용 등 세부 하이퍼파라미터의 민감도와 확장성에 대한 논의가 제한적으로 보임
• 다중 목표 수가 매우 많아질 경우 Pareto Frontier 계산 및 self-contrastive 신호 도출의 계산 복잡도가 어떻게 확장되는지 추가 검증이 필요함
• 본 방법이 특정 diffusion 기반 protein 모델(backbone)에 의존하는지, 다른 생성 모델 계열로의 일반화 가능성에 대한 추가 실험이 요구됨
• 후속 연구로는 memory bank의 장기 축적에 따른 편향 문제, 그리고 실제 wet-lab 검증을 통한 생물학적 타당성 확인이 필요할 것으로 보임

Evaluation

총평: Memory bank와 self-contrastive learning, Pareto alignment를 결합해 test-time diffusion 기반 다중 목표 단백질 설계의 핵심 난제들을 체계적으로 해결한 독창적이고 실용적인 연구로, 향후 다양한 생성 모델 및 응용 분야로의 확장이 기대된다.