LeFlur: A Biomolecular Design Model with Latent Structure Tokens

Motivation

• Known: 기존 protein design 파이프라인은 backbone generation(RFdiffusion 등), sequence design(ProteinMPNN 등), structure prediction(AlphaFold 등)을 각각 별도의 특화 모델로 처리하며, ESM3, DPLM-2 등이 구조를 discrete token으로 변환해 단일 모델에 통합하려는 시도를 해왔다.
• Gap: ESM3, DPLM-2 같은 기존 discrete structure token 모델은 equivariant 구조를 처리하기 위해 복잡한 protein-specific 아키텍처와 큰 codebook(4096~8192 토큰)이 필요하고, protein-only 시스템에 국한되어 ligand-conditioned design을 지원하지 못한다는 근본적 한계가 있다.
• Why: 단백질뿐 아니라 소분자(ligand)까지 포함하는 통합 구조 토큰화가 가능해지면 protein-ligand 복합체 설계, inverse/forward folding, unconditional generation을 단일 경량 모델로 처리할 수 있어 design campaign의 계산 비용과 파이프라인 복잡도를 크게 줄일 수 있다.
• Approach: LatentGenerator라는 표준 Vision Transformer autoencoder와 Gumbel-Softmax 기반 Simple Linear Quantizer(SLQ)로 3D 좌표를 최소 vocabulary discrete token으로 변환하고, 이를 NeoBERT 기반 text transformer에 넣어 discrete flow matching으로 학습한다.

Achievement

1. LatentGenerator를 통한 modality-agnostic 구조 토큰화: equivariant layer 없이 SE(3) data augmentation만으로 학습한 vanilla ViT autoencoder가 단백질, 소분자, protein-ligand 복합체를 하나의 encoder-decoder로 discretize하며, single-step decoding만으로 충분함을 보였다.
2. Simple Linear Quantization(SLQ) 제안: 학습된 linear projection과 Gumbel-Softmax categorical bottleneck만으로 단백질 backbone geometry를 K=256의 극소 codebook으로 압축, DPLM-2(8192)·ESM3(4096) 대비 수십 배 작은 vocabulary로 고품질 geometry 복원을 달성했다.
3. LeFlur 통합 설계 모델: LatentGenerator 토큰을 NeoBERT에 직접 연결하고 discrete flow matching으로 학습해, 단일 weight set으로 inverse folding, forward folding, unconditional generation을 protein-only 및 protein-ligand 시스템 모두에서 특화 baseline과 경쟁력 있는 성능으로 해결했다.
4. 내장형 inference-time verifier: 모델 자체의 pseudo-likelihood(PLL)를 best-of-N ranker로 활용하고, 외부 scorer나 재학습 없이 forward-fold consistency 기반 Self-Reflection 루프로 designability를 향상시켰다.

How

• LatentGenerator: ViT encoder Eθ가 3D 좌표를 continuous embedding으로 매핑하고, quantizer Qϕ(SLQ 또는 protein-ligand 복합체용 FSQ)가 이를 discrete token으로 변환하며, ViT decoder Dψ가 좌표를 재구성함.
• SLQ: continuous embedding Z에 LayerNorm과 학습된 projection matrix Wq를 적용해 logit을 계산하고, Gumbel-Softmax estimator(온도 τ=0.5)로 미분 가능한 discrete sampling을 수행하여 K=256 codebook을 사용.
• protein-ligand 복합체에서는 SLQ가 성능 저하를 보여 Finite Scalar Quantization(FSQ)으로 전환함.
• LeFlur: LatentGenerator의 discrete token(단백질용 Q^(p), ligand용 Q^(ℓ))을 NeoBERT 인코더에 넣어 서열과 구조 토큰을 동일 attention stream에서 처리, Discrete Flow Matching으로 서열/구조 timestep을 독립적으로 샘플링해 학습.
• masked any-order 학습 방식 덕분에 sequence/structure/joint pseudo-likelihood(PLL)를 계산할 수 있으며 이를 best-of-N ranking에 활용.
• Self-Reflection: 생성된 서열을 동일 모델로 재폴딩하여 원래 샘플된 구조와 일치할 때만 채택하는 반복적 refinement loop를 inference 시점에 적용.

Originality

• 단백질과 소분자를 하나의 동일한 shared-vocabulary discrete token 체계로 통합 처리한 최초의 structure-token 모델이라는 점에서 독창적임.
• 복잡한 equivariant 아키텍처 대신 표준 ViT와 SE(3) augmentation만으로 구조 토큰화를 수행하는 단순화 접근을 채택.
• Gumbel-Softmax 기반의 Simple Linear Quantizer(SLQ)를 기존 FSQ 대비 극소 codebook(K=256)으로 적용해 minimal-vocabulary 구조 토큰화를 실현.
• 별도 외부 scorer 없이 모델 자체의 pseudo-likelihood와 forward-fold consistency만으로 inference-time verification(Self-Reflection)을 수행하는 self-contained 설계.

Limitation & Further Study

• SLQ가 protein-ligand 복합체 재구성에서는 성능이 급격히 저하되어 FSQ로 대체해야 하는 한계가 있으며, 이는 소분자와 단백질 간 통합 tokenization의 완전성에 의문을 남김.
• 논문 발췌본에서는 정량적 벤치마크 결과(designability, RMSD 등 구체적 수치)가 충분히 제시되지 않아 baseline 대비 실제 성능 우위 정도를 판단하기 어려움.
• Self-Reflection이 forward-fold consistency에만 의존하므로, 실제 기능적/물리화학적 유효성까지 보장하는지는 추가 검증이 필요함.
• 후속 연구로는 더 다양한 ligand 화학 공간 및 복합체 규모에 대한 tokenization 확장성 검증, 그리고 SLQ와 FSQ를 통합하거나 대체할 수 있는 quantization 기법 탐색이 필요해 보임.

Evaluation

총평: 단백질과 리간드 구조를 단일한 discrete token 체계로 통합하고 이를 표준 text transformer에 접목시킨 아키텍처적 단순화는 실용적 가치가 크며, self-contained 검증 메커니즘(PLL, Self-Reflection)도 흥미로운 기여이나, protein-ligand 복합체에서의 quantization 한계와 구체적 정량 결과의 추가 검증이 필요하다.