기술 요약: 구조 기반 LLM 추론을 위한 확장성 인식 어댑터 (Cuttlefish)

1. 문제 제기

대규모 언어 모델 (LLM) 은 increasingly 분자, 단백질, 핵산과 같은 2 차원 및 3 차원 과학적 구조에 대한 추론에 적용되고 있습니다. 그러나 기존 방법들은 모든 원자 수준의 일반화된 추론을 저해하는 두 가지 중요한 한계에 직면해 있습니다:

1. 고정 길이 병목 현상 (예산 확장성): 현재 아키텍처는 일반적으로 고정 길이의 쿼리 커넥터 (예: Q-Former 스타일) 나 시퀀스 기반 토큰화를 의존합니다. 이러한 접근 방식은 가변 크기의 구조적 그래프를 일정한 토큰 예산으로 압축합니다. 구조적 복잡성이 증가함에 따라 (예: 더 큰 단백질 또는 핵산), 이는 기하학적 특징의 과도한 압축을 초래하여 성능을 저하시킵니다. 반대로, 모든 입력에 대해 토큰 예산을 단순히 증가시키는 것은 비효율적인 할당으로 이어져 작은 개체에 대한 주의를 희석시킵니다.
2. 구조적 환각: 시퀀스 입력 (예: SMILES, 잔기 시퀀스) 만을 의존하는 모델들은 명시적인 기하학적 근거가 부족합니다. 검증 가능한 3 차원 구조적 증거 없이는 LLM 이 조건부 컨텍스트에 의해 추론되지 않는 기하학적 근거를 생성하는 경향이 있어, 작용기, 메커니즘, 또는 장거리 상호작용에 대한 환각을 초래합니다.

본 논문은 이질적인 모든 원자 개체에 대한 통합 인터페이스가 부재하다고 주장하며, 구조적 증거가 어떻게 선택되고 언어 모델에 노출되는지에 대한 재고를 필요로 한다고 역설합니다.

2. 방법론: Cuttlefish

저자들은 언어 추론을 기하학적 단서에 기반하게 하고 구조적 복잡성에 따라 모달리티 토큰을 확장하도록 설계된 통합 멀티모달 LLM 인 Cuttlefish를 소개합니다. 아키텍처는 세 가지 주요 구성 요소로 이루어져 있습니다:

A. 인코더 (SE(3)-공변 EGNN)

입력 원자 좌표와 특징은 먼저 **공변 그래프 신경망 (EGNN)**을 사용하여 인코딩됩니다. 이는 표현이 회전과 병진에 대해 불변임을 보장하여 원자 수준의 모달리티 임베딩을 생성합니다. 인코더는 공간적 관계를 학습하기 위해 마스크된 재구성 작업 (원소 식별, 쌍별 거리, 방향 노이즈 예측) 을 통해 사전 훈련됩니다.

B. 확장성 인식 패치화

고정 예산 병목 현상을 해결하기 위해 Cuttlefish 는 지시 조건부 게이트 메커니즘을 사용합니다:

• 질량 기반 앵커 선택: 지시 임베딩과 노드 임베딩이 주어지면, 스코어링 게이트가 노드별 관련성 로짓을 계산합니다. 앵커는 누적 질량 임계값 ($\rho$) 을 기반으로 적응적으로 선택됩니다. 이는 쿼리 토큰 수 ($k_g$) 가 개체의 구조적 복잡성에 따라 확장되도록 하여, 가장 지시 관련성이 높은 원자만 선택합니다.
• 소프트 패치 성장: 선택된 앵커는 소프트 할당 메커니즘을 통해 구조적 패치로 확장됩니다. 이는 공간적 근접성 (유클리드 거리) 과 의미적 편향 (지시 관련성) 을 융합합니다.
• 패치 내 가중 풀링: 각 패치 내의 노드 특징은 멤버십 가중치를 사용하여 집계되어 구조적 쿼리 토큰을 생성합니다. 이를 통해 모델은 복잡한 영역에 대한 세밀한 기하학적 신호를 보존하면서도 덜 관련 있는 영역을 압축할 수 있습니다.

C. 기하학적 근거 어댑터

구조적 환각을 완화하기 위해 적응형 토큰은 정제되어 LLM 에 주입됩니다:

• 크로스 어텐션 검색: 패치 토큰은 일련의 퓨전 블록을 통해 전체 모달리티 임베딩 (모든 원자) 과 크로스 어텐션을 수행합니다. 이는 쿼리가 패치화 단계에서 추상화된 고해상도 기하학적 단서를 검색할 수 있게 합니다.
• 주입: 정제된 모달리티 토큰은 LLM 의 임베딩 공간으로 투영되어 지시 시퀀스 내 지정된 플레이스홀더 위치에 삽입됩니다. 이는 명시적인 기하학적 단서를 LLM 에 노출시켜, 검증 가능한 구조에 기반한 출력을 조건부로 설정합니다.

D. 훈련 프로토콜

모델은 세 단계로 훈련됩니다:

1. 인코더 사전 훈련: EGNN 은 모든 원자 공간 그래프에서 사전 훈련됩니다.
2. 모달리티 정렬 튜닝: EGNN 과 LLM 은 고정되며, 커넥터 (패치화 및 어댑터) 만 저자들이 소개한 새로운 모든 원자 지시 데이터셋인 GEO-AT에서 최적화됩니다.
3. LLM 적응 튜닝: LLM 은 새로운 모달리티 주입 분포를 수용하고 일반적인 언어 능력을 보존하기 위해 더 작은 학습률로 고정 해제됩니다.

3. 주요 기여

1. 확장성 인식 패치화: 모달리티 토큰 수가 구조적 복잡성에 따라 확장되도록 보장하는 지시 조건부 앵커 패치 성장 메커니즘으로, 고정 길이 커넥터 병목 현상을 효과적으로 완화합니다.
2. 기하학적 근거 어댑터: 모달리티 토큰을 통해 검증 가능한 기하학적 단서를 LLM 에 주입하여, 기하학에 명시적으로 조건을 부여함으로써 구조적 환각을 줄이는 모듈입니다.
3. Cuttlefish 프레임워크: 이질적인 모든 원자 모달리티 (분자, 단백질, DNA, RNA) 에 대한 추론이 가능하며 벤치마크 전반에 걸쳐 강력한 성능을 보이는 통합 구조 인식 LLM 입니다.
4. GEO-AT 데이터셋: 원자 수준 좌표와 언어 주석이 포함된 분자, 단백질, 핵산을 포괄하는 최초의 모든 원자 지시 데이터셋을 소개합니다.

4. 실험 결과

실험은 Mol-Instructions, DNA-Chat, RNA-QA를 포함한 이질적인 모든 원자 벤치마크에서 수행되었습니다.

• 일반 LLM 대비 성능: Cuttlefish 는 모든 모달리티에 걸쳐 일반 LLM 베이스라인 (예: Llama-3, Qwen, Mistral) 및 추론 강화 변형 모델을 크게 능가합니다. 예를 들어, GEO-AT 벤치마크에서 Cuttlefish + Llama-3.1-8B 는 평균 METEOR 점수 0.428과 BERTScore 0.864를 기록한 반면, 백본 단독은 각각 0.186과 0.694를 기록했습니다.
• 모달리티별 베이스라인 대비 성능: Cuttlefish 는 Mol-Llama (분자), ProtChatGPT (단백질), ChatNT (핵산) 와 같은 전문 모델을 능가하여 대부분의 작업에서 1 위 성능을 달성합니다.
• 확장성 안정성: 개체 크기가 증가함에 따라 성능이 저하되는 고정 길이 커넥터와 달리, Cuttlefish 는 다양한 그래프 노드 수 (작은 분자에서 큰 단백질까지) 에 걸쳐 안정적인 생성 품질을 유지하며 유리한 확장 특성을 보여줍니다.
• 토큰 효율성: Cuttlefish 는 성능을 유지하면서 선형 확장 베이스라인 (예: ChatNT) 에 비해 훨씬 낮은 구조적 토큰 예산을 유지하여 적응형 토큰 할당의 효율성을 확인합니다.
• 환각 감소: 환각 평가 (작용기 예측, 메커니즘 설명, 장거리 상호작용) 에서 Cuttlefish 는 가장 낮은 환각률 (예: 단백질 작용기의 경우 0.04, ProtChatGPT 는 0.10) 과 가장 높은 답변률을 달성합니다.
• 강건성: 모델은 좌표 노이즈 하에서 우아하게 저하되며, 실험적 구조가 AlphaFold2 예측으로 대체되더라도 여전히 효과적입니다.

5. 중요성과 주장

본 논문은 Cuttlefish가 모든 원자 모달리티에 걸친 구조 기반 추론을 위한 통합 프레임워크를 확립한다고 주장합니다. 그 중요성은 확장성과 신뢰성이라는 이중적 과제를 해결하는 데 있습니다:

• 확장성: 경직된 고정 길이 커넥터를 확장성 인식 패치화로 대체함으로써 아키텍처는 정보 병목 현상을 우회하여 과도한 압축이나 컨텍스트 비효율성 없이 구조적 복잡성을 처리할 수 있게 합니다.
• 신뢰성: 기하학적 근거 어댑터를 사용하여 명시적인 기하학적 증거를 주입함으로써, 시퀀스-only 입력과 관련된 불가피한 불확실성 (베이즈 위험) 을 줄여 구조적 환각을 억제합니다.

저자들은 시퀀스-only 입력이 구조 조건부 이해에 불충분하며, 고정 길이 커넥터가 기하학 의존적 의미론을 포착하는 데 근본적인 병목 현상을 부과한다고 결론지었습니다. 그들은 적응형 토큰 확장을 통한 결합 모달리티 훈련이 견고하고 통합된 멀티모달 과학적 추론의 핵심이라고 주장합니다. 이 연구는 일반적인 과학적 모델링 환경을 대상으로 하며, 해당 연구 분야에 공통적인 것 이상의 광범위한 사회적 영향을 목표로 하지 않습니다.