DiagramBank: 검색 증강 생성을 위한 논문 메타데이터 포함 대규모 다이어그램 디자인 예시 데이터셋

자율적인 AI 과학자 시스템이 텍스트와 코드는 잘 작성하지만, 논문의 핵심인 '티저 다이어그램' 생성에는 여전히 어려움을 겪고 있다. 이 논문은 고품질 다이어그램 예시와 풍부한 텍스트 문맥을 결합한 데이터셋을 제공하여, AI가 전문가 수준의 학술 도식을 생성할 수 있는 기반을 마련했다.

기존의 텍스트-이미지 생성 모델은 복잡한 논리 흐름과 텍스트가 결합된 학술 다이어그램을 생성할 때 화살표 연결이 어색하거나 텍스트가 깨지는 등 구조적 일관성이 부족한 한계가 있다. 이는 모델이 학술적 도식의 시각적 관습(Visual Convention)을 충분히 학습하지 못했기 때문이다.

이 문제를 해결하기 위해 연구진은 '검색 증강 생성(RAG)' 개념을 이미지 생성에 도입했다. 모델이 무에서 유를 창조하는 대신, DiagramBank에서 검색된 고품질의 실제 논문 도식을 '시각적 가이드(Visual Prior)'로 참고하게 함으로써 레이아웃, 색상 조합, 아이콘 스타일을 모방하도록 유도한다.

결과적으로 모델은 검색된 예시의 구조적 특징을 유지하면서도 사용자가 입력한 새로운 논문 내용에 맞춰 세부 요소를 변경할 수 있게 된다. 이는 마치 숙련된 디자이너가 기존의 잘 만들어진 템플릿을 참고하여 새로운 도식을 그리는 과정과 유사하게 동작하여 생성물의 전문성을 크게 높인다.

전체 시스템은 데이터 수집, 추출, 분류, RAG 파이프라인의 4단계로 구성된다. OpenReview에서 수집된 PDF를 PDFFigures 2.0으로 파싱하여 이미지와 캡션을 추출하고, PyMuPDF를 통해 해당 이미지가 언급된 본문 문맥을 확보한다.

핵심 분류 메커니즘은 CLIP(ViT-B-32)을 활용한다. [이미지 임베딩 → 텍스트 임베딩(diagram, plot, photo, other)과의 코사인 유사도 계산 → Softmax 적용] 과정을 거쳐 0.85 이상의 신뢰도를 가진 다이어그램만 선별한다. 이 과정에서 단순 차트와 복잡한 도식을 효과적으로 분리한다.

RAG 파이프라인은 계층적 검색을 수행한다. 먼저 논문 제목 인덱스에서 후보군 S1을 추출하고, S1 내에서 초록 인덱스를 통해 방법론이 유사한 S2를 재정렬(Rerank)한다. 마지막으로 캡션 인덱스를 통해 세부 구성 요소가 일치하는 최종 예시를 선택하여 생성 모델의 프롬프트에 주입한다.

DiagramBank는 2017년부터 2025년까지 ICLR, ICML, NeurIPS, TMLR에서 발표된 논문들을 포함하며, 최종적으로 89,422개의 다이어그램을 확보했다. 분석 결과, 최근 AI 연구의 폭발적 증가와 함께 다이어그램의 양도 기하급수적으로 증가하고 있으며, 캡션의 길이는 점차 짧아지는 추세를 보였다.

사례 연구(Case Study)에서 RAG를 적용하지 않은 기본 모델은 원색 위주의 단순한 레이아웃을 생성한 반면, DiagramBank-RAG를 적용한 모델은 검색된 예시의 파스텔톤 색상, 둥근 모서리 상자, 순환형 워크플로우 구조를 성공적으로 전이(Transfer)하여 실제 논문에 바로 사용할 수 있는 수준의 품질을 보여주었다.

데이터셋은 DuckDB를 사용하여 관계형으로 관리되며, 각 레코드는 논문 수준의 메타데이터와 피규어 수준의 문맥이 비정규화(Denormalized)된 형태로 저장되어 검색 효율을 극대화했다. 임베딩 모델로는 OpenAI의 text-embedding-3-large를 사용하며, FAISS 라이브러리를 통해 고속 근사 최근접 이웃(ANN) 검색을 지원한다.

특히 'Deep Fetch' 메커니즘을 도입하여, 상위 단계 필터링에서 발생할 수 있는 재현율(Recall) 손실을 방지하기 위해 필터링된 범위보다 더 넓은 범위의 후보를 하위 단계로 전달한 후 정교하게 정렬하는 방식을 채택했다. 이는 단순한 키워드 매칭을 넘어선 심층적인 문맥 일치를 가능하게 한다.

자동화된 휴리스틱에 의존하므로 캡션이 불완전하거나 본문 인용 문맥이 정확하지 않은 노이즈가 포함될 수 있다. 또한 현재 이미지 생성 모델의 한계로 인해 아주 복잡한 화살표 위상이나 미세한 텍스트 렌더링은 여전히 추가적인 편집이 필요할 수 있다.