ASK2024

DDPM 기반 데이터 증강과 준지도학습을 결합해 OCSR(광학 분자구조 인식) 성능 향상 연구

Note: This article was drafted with AI assistance and reviewed by Jin Hyuk Kim.

김진혁 (Jin Hyuk Kim)1, 송태웅 (Tae Woong Song)1, 최종환 (Jonghwan Choi)2

1한림대학교 소프트웨어학부 · 2한림대학교 정보과학대학

Venue: 한국정보처리학회(KIPS) 춘계학술대회 ASK2024 · 서울대학교 평창캠퍼스 · 2024

Paper · Slides

Summary

DDPM[1] 기반 분자구조 이미지 생성과 semi-supervised learning(pseudo-labeling)을 결합해, 학습 데이터가 부족한 상황에서도 OCSR(Optical Chemical Structure Recognition) 모델의 인식 정확도를 5.4%p 끌어올린 연구입니다.

Background: 스타일이 다른 분자구조 이미지엔 왜 약할까?

OCSR(Optical Chemical Structure Recognition)은 문헌 속 분자구조 이미지를 인식해 SMILES[2] 문자열로 변환하는 기술로, 화학정보학 데이터 수집의 핵심 도구입니다. 하지만 MolScribe[3] 등 기존 OCSR 모델은 PubChem[5]처럼 대규모 데이터로 pretrained 되어 있어, ChemDraw[3]로 그린 것처럼 스타일이 다른 이미지에는 인식 정확도가 크게 떨어집니다.

Bottleneck: 스타일별로 fine-tuning 하려 해도, 다른 스타일의 소규모 데이터만으로는 충분한 학습이 어렵습니다. 다량의 labeled dataset(이미지-SMILES 쌍)을 수작업으로 만드는 데는 상당한 시간과 비용이 들기 때문입니다.

Method: DDPM 기반 Data Augmentation과 Pseudo-labeling

소규모 데이터의 한계를 극복하기 위해, generative model(DDPM)로 부족한 이미지를 만들고 semi-supervised learning으로 라벨을 채우는 4단계 fine-tuning 전략을 제안했습니다.

  1. DDPM train: ChemDraw 스타일 분자구조 이미지로 DDPM(Denoising Diffusion Probabilistic Model)을 학습
  2. DDPM sampling: 학습된 DDPM으로 새로운 분자구조 이미지를 대량 생성
  3. OCSR inference (Pseudo-labeling): 생성된 이미지에 pretrained OCSR 모델을 적용해 SMILES pseudo label을 예측
  4. OCSR retrain: ChemDraw 원본 데이터 + DDPM 생성 데이터를 합쳐 OCSR 모델을 retrain(fine-tuning)

DDPM은 이미지만 생성할 뿐 대응되는 SMILES 라벨이 없다는 문제가 있는데, 이를 pseudo-labeling으로 보완하는 것이 핵심 아이디어입니다.

Results: 인식 정확도 향상

정답 SMILES와 예측 SMILES의 유사도를 Tanimoto Similarity[4]로 측정해, ChemDraw 데이터셋(train 4,163개 / test 500개)에 대해 교차검증으로 성능을 비교했습니다. Fine-tuning에 사용한 데이터 크기가 클수록 (원본 ChemDraw만 사용 → DDPM 생성 데이터까지 추가) 성능이 함께 향상되는 것을 확인했습니다.

Method Dataset size Tanimoto Similarity
Pretrained (w/o fine-tuning) 0 71.4%
ChemDraw fine-tuning 4,163 84.7%
ChemDraw + DDPM fine-tuning (ours) 14,728 90.1%

DDPM으로 생성한 데이터를 함께 사용했을 때, ChemDraw만 사용한 fine-tuning 대비 5.4%p의 인식 정확도 향상(84.7% → 90.1%)을 확인했습니다.

Conclusion

소규모 데이터셋만으로는 OCSR 모델을 특정 이미지 스타일에 fine-tuning하기 어렵다는 문제를, DDPM 기반 data augmentation과 pseudo-labeling을 결합해 완화할 수 있음을 보였습니다. 제안한 학습 전략은 baseline 대비 5.4%p의 성능 향상을 가져왔으며, 이 연구는 한국정보처리학회 춘계학술대회 ASK2024 학부생 세션에서 발표되었습니다.

이 연구는 이후 연구 고도화를 거쳐, 정형화된 분자구조 이미지뿐 아니라 hand-drawn 분자구조 이미지까지 다루도록 확장된 OCSAug 연구로 이어졌습니다.

References

  1. Nichol, Alexander Quinn, and Prafulla Dhariwal. “Improved Denoising Diffusion Probabilistic Models.” International Conference on Machine Learning. PMLR, 2021.
  2. Weininger, David. “SMILES, a Chemical Language and Information System. 1. Introduction to Methodology and Encoding Rules.” Journal of Chemical Information and Computer Sciences 28.1 (1988): 31-36.
  3. Qian, Yujie, et al. “MolScribe: Robust Molecular Structure Recognition with Image-to-Graph Generation.” Journal of Chemical Information and Modeling 63.7 (2023): 1925-1934.
  4. Bajusz, Dávid, Anita Rácz, and Károly Héberger. “Why is Tanimoto Index an Appropriate Choice for Fingerprint-Based Similarity Calculations?” Journal of Chemical Information and Modeling 50.5 (2010): 742-754.
  5. Kim, Sunghwan, et al. “PubChem 2019 Update: Improved Access to Chemical Data.” Nucleic Acids Research 47.D1 (2019): D1102-D1109.