Empirical Observations on Parameter Scaling in Chemical Language Models

Motivation

• Known: 그래프 기반 아키텍처(GNN)에서는 파라미터 수 증가가 분자 특성 예측 성능을 향상시키는 scaling 행동이 관찰되었고, 자연어 처리 분야에서는 Kaplan et al., Hoffmann et al. 등을 통해 neural scaling law가 잘 정립되어 있다.
• Gap: SMILES 기반 autoregressive chemical language model에서도 NLP나 GNN에서 관찰된 scaling law 원칙이 유사하게 전이되는지에 대한 연구는 아직 충분히 이루어지지 않았다.
• Why: 화학 언어 모델의 규모 확장이 실제로 유의미한 성능 향상을 보장하는지 확인하는 것은 향후 대형 화학 foundation model 개발 및 자원 투자 방향성 결정에 중요한 시사점을 제공한다.
• Approach: Chemlactica 계열의 170M, 380M, 1.3B 파라미터 LLaMA 기반 모델을 동일한 PubChem 사전학습 코퍼스로부터 시작하여, Polaris ADME multi-task regression, PXR induction single-task regression, BELKA binding classification 세 가지 downstream task에 대해 fine-tuning 후 성능을 비교한다.

Achievement

Figure 1. Metric vs Model size across evaluated tasks.

1. PXR induction task에서 일관된 스케일링 관찰: 모델 크기가 170M에서 1.3B로 증가함에 따라 예측 성능이 단조롭게 향상되는 가장 명확한 scaling 행동을 확인함.
2. Polaris ADME에서 부분적 스케일링 효과: 1.3B 모델이 170M 대비 R2에서 약 20% 상대적 향상을 보였고 대부분의 endpoint에서 MSE/MAE가 감소했으나, task별로 일관되지 않은 변동성을 보임.
3. BELKA 데이터셋에서 일반화 실패 관찰: 모델 크기 증가가 public test set에서는 성능 향상을 가져왔으나, out-of-distribution 성격이 강한 private test set에서는 일관되게 일반화되지 않음을 확인함.
4. fine-tuning 전략(layer unfreezing, pooling, loss function 등)이 스케일링 효과에 미치는 영향을 체계적으로 분석: task별 최적 fine-tuning 구성(layer unfreezing 비율, pooling 방식, Focal Loss 등)이 스케일링 이득의 발현 여부에 중요한 역할을 함을 보임.

How

Figure 1. Metric vs Model size across evaluated tasks.

• 모델: Chemlactica (LLaMA 3 기반) 170M, 380M, 1.3B 파라미터 변형, 모두 동일한 약 1.1억 개 분자로 구성된 PubChem corpus(40B tokens)로 사전학습됨.
• Fine-tuning 전략: 사전학습된 backbone을 우선 전체 동결한 후, catastrophic forgetting을 완화하기 위해 마지막 transformer layer들과 normalization layer만 선택적으로 unfreeze함. Bayesian search를 통해 하이퍼파라미터 최적화 수행.
  • Polaris: 모든 변형에서 마지막 2개 layer unfreeze (~10% 파라미터 학습)
  • BELKA: 파라미터 비율 31-34% 유지하도록 170M/380M/1.3B에 각각 6/8/7 layer unfreeze
  • PXR induction: 모델별 독립적으로 튜닝하여 3/10/10 layer unfreeze
• Differential learning rate: backbone과 classification head를 서로 다른 learning rate 그룹으로 분리하여 optimizer에 전달 (backbone lr이 head보다 1~3 orders of magnitude 작음).
• Pooling: 회귀 task(Polaris, PXR)에는 mean pooling, BELKA 분류 task에는 last-token pooling([END SMILES] 토큰) 사용.
• Loss function: BELKA의 극심한 class imbalance(양성 샘플 0.5% 미만) 문제를 해결하기 위해 Focal Loss(γ=2.0)와 label smoothing을 결합하여 사용.
• Optimizer/Scheduler: AdamW, linear warmup 후 cosine annealing.
• BELKA data split: 데이터 누수 방지 및 진정한 out-of-distribution 일반화 평가를 위해 building block 기준 scaffold splitting(3% building block held-out) 적용.

Originality

• 기존 scaling law 연구가 주로 NLP 또는 GNN 기반 분자 모델에 집중된 것과 달리, autoregressive SMILES 기반 chemical language model에서의 파라미터 스케일링 효과를 실증적으로 검증한 최초의 시도 중 하나임.
• 단일 벤치마크가 아닌 회귀(Polaris, PXR)와 분류(BELKA), 소규모 및 대규모 데이터셋을 모두 포함하여 task 특성에 따른 스케일링 효과의 이질성을 다각도로 조명함.
• BELKA에서 building block 기반 scaffold split을 통해 public/private test 간 실질적인 out-of-distribution 일반화 격차를 드러낸 점이 특징적임.

Limitation & Further Study

• 평가된 모델 규모가 170M~1.3B로 제한적이며, 진정한 scaling law를 도출하기에는 규모 범위와 모델 variant 수가 부족함(저자들도 명시적으로 인정).
• 각 downstream task마다 서로 다른 fine-tuning 전략(layer unfreezing 비율, pooling 등)을 사용하여, 관찰된 scaling 효과가 모델 자체의 특성인지 fine-tuning 설정의 차이인지 완전히 분리하기 어려움.
• Polaris, PXR 데이터셋은 표본 수가 3-4.5천 개로 작아 통계적 유의성 및 재현성 검증이 제한적일 수 있음.
• 후속 연구로 더 다양한 크기와 아키텍처의 chemical language model을 포함한 체계적 scaling law 정립이 필요함.

Evaluation

총평: 화학 언어 모델의 스케일링 행동에 대한 실증적이고 체계적인 초기 탐색으로서 의미 있는 워크숍 수준의 기여를 하지만, 모델 규모 범위가 제한적이고 task별 fine-tuning 전략 차이로 인해 명확한 scaling law 결론을 내리기에는 한계가 있다.