기술 요약: SMolLM - 소형 언어 모델이 소형 분자 문법을 학습하다

문제 제기

새로운 분자를 설계하는 것은 신약 개발 및 재료 과학에서 중요한 과제이며, SMILES(단순화된 분자 입력 라인 엔트리 시스템) 문자열을 생성하는 언어 모델 (LM) 에 의해 점점 더 해결되고 있습니다. 최근의 분자 생성기들은 수억 개의 파라미터로 확장되었지만, 화학적으로 유효한 구조를 생성할 수 있게 하는 근본적인 계산 메커니즘은 여전히 잘 이해되지 않고 있습니다. 구체적으로, 그러한 거대한 파라미터 수가 필요한지, 또는 어떤 특정 계산이 모델로 하여금 SMILES 의 엄격한 구문 및 화학적 제약 조건 (예: 괄호 일치, 고리 숫자 짝짓기, 유효한 원자 원자가 유지) 을 충족하게 하는지는 명확하지 않습니다. 또한, 각 레이어가 고유한 가중치를 갖는 표준 트랜스포머 아키텍처에서는 레이어 간의 기능적 차이가 계산적 진보에서 비롯된 것인지, 아니면 단순히 파라미터 차이에서 비롯된 것인지 구분하기 어렵습니다.

방법론

저자들은 가중치 공유 트랜스포머 아키텍처 기반의 소형 분자 생성기인 SMolLM을 소개합니다. 이 모델은 서로 다른 레이어를 적층하는 대신, 단일 트랜스포머 블록을 $K$번 반복하여 (가상 깊이) 토큰을 자기회귀적으로 생성합니다.

• 아키텍처: 핵심 모델은 현대적 GPT 설계 (회전 위치 임베딩, SwiGLU 활성화, 프리-노름) 기반의 53K 파라미터 가중치 공유 블록 (WS-53K) 입니다. 이 블록은 토큰 생성 단계당 $D=8$번 적용됩니다.
• 학습: 모델들은 문자 단위 토큰화 및 SMILES 증강 ($n_{aug}=10$) 을 사용하여 ZINC-250K 데이터셋 (249,455 개의 약물 유사 분자) 으로 학습되었습니다. 학습에는 100 에포크에 걸쳐 코사인 학습률 스케줄을 사용한 AdamW 가 활용되었습니다.
• 평가: 유효성은 RDKit 을 통해 실제 분자로 디코딩 가능한 생성된 문자열의 비율로 측정되었습니다. 추가 지표로는 FCD(프레셰트 케민넷 거리), 내부 다양성, 고유성, 그리고 신규성이 포함되었습니다.
• 메커니즘 분석: 모델이 어떻게 학습하는지 이해하기 위해 저자들은 반복적 패스 전반에 걸쳐 세 가지 보완적 분석을 수행했습니다:
  1. 오류 분류: 다른 가상 깊이에서 실패 모드 (괄호 불일치, 고리 숫자 불일치, 원자가 위반) 에 따라 유효하지 않은 분자를 분류합니다.
  2. 표현 프로빙: 선형 프로브와 희소 오토인코더 (SAE) 를 사용하여 각 패스의 숨겨진 상태 내에서 화학적 특징 (괄호 깊이, 고리 상태, 원자 식별) 을 탐지합니다.
  3. 인과적 제거: 특정 패스에서 특정 어텐션 헤드를 0 으로 설정하여 오류율에 미치는 인과적 영향을 관찰합니다.

주요 결과

1. 파라미터 효율성 및 유효성

53K 파라미터 가중치 공유 모델 (WS-53K) 은 ZINC-250K 벤치마크에서 95.3% 의 유효성을 달성합니다.

• 비교: 이는 527K 파라미터 (10 배 더 큼) 를 가진 동일하게 학습된 비공유 표준 GPT(87.6% 유효성 달성) 보다 우수합니다.
• 파레토 프론티어: 가중치 공유 모델은 표준 GPT 들에 비해 훨씬 적은 파라미터로 더 높은 유효성을 제공하며, 서브 메가파라미터 파레토 프론티어를 지배합니다.
• 강건성: 성능 향상은 다양한 샘플링 온도, 톱-k 설정, 그리고 매칭된 파라미터 수의 순환적 베이스라인 (GRU) 에 대해 유지됩니다. 지식 증류 및 직접 선호도 최적화 (DPO) 와 같은 사후 학습 방법은 유효성을 저하시키는 것으로 밝혀졌으며, 이는 성능이 정렬 기법이 아닌 반복적 계산 자체에서 비롯됨을 시사합니다.

2. 반복적 계산과 계층 구조

이 연구는 공유 블록이 모든 제약을 동시에 해결하는 대신, 패스를 거쳐 고정된 계층적 순서로 SMILES 제약을 해결함을 보여줍니다:

• 패스 2: 괄호 일치 (로컬 구문) 가 대부분 해결됩니다.
• 패스 4: 고리 닫기 (장거리 숫자 짝짓기) 가 해결됩니다.
• 패스 8: 원자 원자가 (화학적 일관성) 가 해결됩니다.

이 계층 구조는 다음을 통해 확인되었습니다:

• 오류 분류: 패스 2 에 괄호 오류가, 패스 4 에 고리 오류가, 패스 8 에 원자가 오류가 각각 급격히 감소했습니다.
• 표현 분석: 선형 프로브와 SAE 는 초기 패스 (12) 에서 괄호 깊이 특징을, 중간 패스 (46) 에서 고리 상태 특징을, 후기 패스에서 원자 식별 특징을 탐지했습니다.
• 인과적 제거: 패스 1의 특정 어텐션 헤드를 제거한 결과, 고리나 원자가 오류에는 큰 영향을 주지 않으면서 괄호 오류가 22% 포인트 증가하여, 괄호 일치 계산이 첫 번째 패스의 단일 구성 요소에 국한됨을 확인했습니다.

3. 용량 대 깊이

53K 모델과 206K 가중치 공유 모델 (WS-206K) 을 비교한 결과, 더 큰 모델이 제약을 더 빠르게 해결하지만 (예: 작은 모델의 패스 4 대비 패스 2 에서 괄호 일치 해결), 해결 순서는 동일하게 유지됨이 밝혀졌습니다. 용량은 수렴의 속도를 변경할 뿐, 계산 단계의 순서는 변경하지 않습니다.

중요성 및 주장

본 논문은 SMolLM 이 두 가지 주요 목적을 수행한다고 주장합니다:

1. 소형 생성기: 가중치 공유가 고품질 분자 생성에 필요한 파라미터 수를 극적으로 줄일 수 있음을 보여주어, 서브 메가파라미터 모델의 효율성 프론티어를 이동시킵니다.
2. 반복적 계산을 위한 테스트베드: 동일한 가중치를 재사용함으로써 계산의 시간적 진전을 가시화합니다. 고정된 계층 구조 (괄호 $\to$ 고리 $\to$ 원자가) 는 계층적 제약을 가진 형식 언어 도메인이 깊고 서로 다른 레이어를 통해가 아니라 반복적 정제를 통해 해결될 수 있음을 시사합니다.

저자들은 이 작업이 작은 모델이 화학 문법을 학습하는 방식에 대한 메커니즘적 설명을 제공하며, 경험적 성능을 SMILES 의 기호적 구조와 연결한다고 강조합니다. 또한 SMILES 표현에 대한 초점 (인코딩과 무관한 분자 구조가 아님) 과 메커니즘 측정의 유한 표본적 성격을 포함한 한계를 지적합니다. 논문은 가중치 공유가 단순한 압축 기술이 아니라, 정확한 기호적 제약을 가진 도메인에서 구조화된 반복적 계산을 드러내는 방법이라고 결론지었습니다.