Steering Sequence Generation in Protein Language Models through Iterative Lookback Monte Carlo Sampling

Motivation

• Known: pLM은 대규모 진화 데이터로부터 학습되어 높은 확률의 서열이 실험적으로 기능적일 가능성이 높다는 것이 알려져 있으며, 이를 활용해 생성물을 특정 속성 쪽으로 steering하는 시도가 이루어져 왔다.
• Gap: 기존 steering 방법들은 low-temperature sampling, RL 기반 정렬, fine-tuning/retraining 등으로 심각한 diversity collapse를 유발하거나 계산 비용이 큰 재학습을 요구하는 한계가 있어, 생성 품질과 steering 목표를 동시에 만족하면서 다양성을 보존하는 training-free 방법이 부재하다.
• Why: 단백질 설계에서 열안정성 향상이나 특정 wild-type 인접 변이체 탐색처럼 실용적 속성을 겨냥한 생성이 중요한데, 기존 방법들은 다양성-품질 트레이드오프에서 급격한 손실을 겪거나 비용이 크므로, 이를 이론적으로 정당화된 최대 엔트로피 원칙 하에 training-free로 해결하는 것은 실질적 응용 가치가 크다.
• Approach: 생성 품질(E)과 steering potential(S)에 대한 기댓값 제약 하에서 엔트로피를 최대화하는 목표 분포 Q*(a) ∝ exp(-βE(a) - λS(a))를 유도하고, 이를 근사 샘플링하기 위해 autoregressive elongation과 MH refinement를 반복하는 ILMC 알고리즘을 설계했다.

Achievement

Figure 2. Diversity vs. generative quality (sDO on PL, 1024

1. 이론적 정당화: Q(a)가 고정된 생성 품질과 steering 제약 하에서 유일한 최대 엔트로피 분포임을 증명하고(Proposition 2.1), 임의의 다른 분포 Q에 대해 엔트로피가 높을수록 Q와의 KL divergence가 작아짐을 보였다(Proposition 2.2).
2. 다양성-품질 트레이드오프 개선: 동일한 생성 품질(matched generative quality) 조건에서 ILMC가 표준 autoregressive baseline보다 더 다양한 샘플을 생성함을 실험적으로 입증했다.
3. 열안정성 steering 성능: 간단한 steering potential을 사용해 ILMC가 비steering autoregressive sampling 및 compute-matched rejection sampling보다 더 높은 예측 melting temperature를 달성했다.
4. classifier-guided steering 확장: ILMC를 non-decomposable objective인 classifier 기반 steering(예: homodimerization 구조 클래스)에 자연스럽게 확장하여, 순수 autoregressive guidance 대비 다양성에서 우위를 보이면서 target property enrichment는 비슷하게 유지했다.
5. 다중 family 및 다중 모델 검증: Chorismate Mutase, Phage Lysozyme, Response Regulator 등 family-specific pLM뿐 아니라 multi-family 모델인 ProGen3에서도 일관된 성능 개선을 확인했다.

How

Figure 1. Iterative Lookback Monte Carlo (ILMC) sampling framework. The algorithm approximates sampling from the target

• 시퀀스 prefix a≤t에 대해 partial generative energy Et(a≤t) (base model의 누적 negative log-likelihood)와 partial steering potential St(a≤t) (미래 토큰에 의존하지 않고 EOS에서 전체 steering score와 일치하도록 정의됨)를 도입한다.
• Autoregressive Elongation 단계: pLM을 사용해 현재 prefix를 블록 단위로 확장한다.
• MH Refinement 단계: 생성 이력의 최근 세그먼트(lookback window)를 pLM으로 재샘플링하여 변이체(variant)를 제안하고, 이를 βE + λS로 정의된 global potential 기준의 Metropolis-Hastings 기준으로 accept/reject한다.
• 이 elongation과 refinement를 EOS 토큰이 나올 때까지 반복하며, block size, lookback window 크기, MC step 수(Nmc) 등의 하이퍼파라미터로 알고리즘을 제어한다.
• classifier-guided steering의 경우 steering potential S를 classifier score로 대체하여 동일한 프레임워크를 그대로 적용한다.

Originality

• 기존 low-temperature sampling(로짓 스케일링)과 진짜 β-tilted 목표 분포 샘플링을 개념적으로 구분하고, 전자가 후자를 정확히 근사하지 못함을 명확히 지적한 점이 독창적이다.
• steering을 위한 최대 엔트로피 원리를 도입하여, "생성 품질 유지 + steering 목표 달성" 하에서 다양성을 이론적으로 보장하는 프레임워크(Proposition 2.1, 2.2)를 제시한 점이 새롭다.
• autoregressive 생성 모델에 EBM 스타일의 MH refinement를 lookback 방식으로 결합하여 training-free inference-time steering을 가능케 한 알고리즘적 조합이 참신하다.
• decomposable potential뿐 아니라 non-decomposable classifier-guided objective까지 동일 프레임워크로 확장 가능함을 보인 점도 실용적 독창성을 더한다.

Limitation & Further Study

• 발췌된 내용에서는 ILMC의 계산 비용(MH refinement로 인한 추가 forward pass 수)에 대한 정량적 분석이나 base autoregressive sampling 대비 오버헤드가 명확히 제시되지 않아, 실제 배포 시 비용-이득 트레이드오프에 대한 추가 검증이 필요하다.
• steering potential S와 classifier score의 설계가 특정 속성(열안정성, wild-type 근접성, homodimerization)에 국한되어 있어, 더 복잡하거나 다중 목적(multi-objective) steering으로의 일반화 가능성에 대한 추가 연구가 요구된다.
• MH refinement의 수렴성(mixing) 및 lookback window 크기·block size 등 하이퍼파라미터 선택이 최종 성능에 미치는 민감도에 대한 체계적 분석이 본문 발췌에서는 부족해 보인다.
• 검증이 특정 protein family(Chorismate Mutase, Phage Lysozyme, Response Regulator)와 ProGen3에 한정되어 있어, 더 다양한 단백질 계열 및 대규모 실험적 검증(실제 wet-lab validation)이 후속 연구로 필요하다.

Evaluation

총평: 이론적으로 잘 정당화된 최대 엔트로피 프레임워크와 training-free inference-time 알고리즘을 결합해 pLM steering의 diversity-quality trade-off 문제를 실질적으로 개선한 견실한 연구로, 실제 단백질 엔지니어링에 바로 적용 가능한 실용적 가치가 높다.