대규모 단일세포 오믹스 학습을 위한 효율적 데이터 로더 scDataset

초록

scDataset는 단일세포 RNA‑seq와 같은 수백만 셀 규모 데이터셋을 디스크에 그대로 두고도 PyTorch와 원활히 연동되는 데이터 로더이다. 블록 샘플링과 배치형 페칭을 결합해 “준무작위” 샘플링을 구현함으로써 순차 스트리밍의 고속 I/O와 무작위 샘플링이 요구하는 미니배치 다양성을 동시에 만족한다. Tahoe‑100M(1억 셀) 실험에서 기존 무작위 접근보다 100배 이상 빠른 처리량을 보였으며, 분류 정확도는 진정한 무작위 샘플링과 차이가 없었다.

상세 분석

본 논문은 대규모 단일세포 오믹스 데이터셋을 딥러닝에 적용할 때 가장 큰 병목인 디스크 I/O 문제를 근본적으로 재설계한다. 기존 접근법은 크게 세 가지로 나뉜다. 첫째, 완전 무작위 샘플링은 셀당 하나씩 랜덤 액세스를 요구해 초당 수십 개 샘플에 머무른다. 둘째, 전체 데이터를 메모리로 로드하면 I/O는 사라지지만, 300 GB 규모의 희소 행렬을 밀집 형태로 변환하면 1 TB 이상이 필요해 일반 워크스테이션에서는 불가능하다. 셋째, 순차 스트리밍은 디스크 연속 읽기로 높은 처리량을 얻지만, 실험 플레이트 단위로 정렬된 데이터가 연속적으로 공급돼 클래스 불균형 및 catastrophic forgetting을 초래한다.

scDataset는 이러한 트레이드오프를 “블록 샘플링”과 “배치형 페칭”이라는 두 단계로 해결한다. 데이터셋을 크기 b인 연속 블록으로 분할하고, 미니배치 크기 m에 대해 ⌈m/b⌉개의 블록을 무작위로 선택한다. 각 블록은 연속 읽기로 한 번에 b개의 셀을 가져오므로 랜덤 디스크 접근 횟수가 m에서 ⌈m/b⌉로 감소한다. 블록 크기가 클수록 동일 블록 내 셀 간 상관성이 높아 미니배치 다양성이 감소하지만, 이는 “배치형 페칭”으로 보완한다. 페칭 팩터 f를 도입해 m·f개의 셀을 한 번에 읽어 메모리 버퍼에 저장하고, 버퍼 내부에서 완전 무작위 셔플 후 f개의 미니배치로 나눈다. 결과적으로 하나의 미니배치가 40개의 블록(예시)에서 추출된 640개의 셀을 섞어 구성되므로, 블록 샘플링만 사용했을 때보다 훨씬 높은 엔트로피를 유지한다.

이론적으로 저자들은 블록 크기 b와 페칭 팩터 f에 대한 미니배치 엔트로피의 기대값 상한을 도출하고, 엔트로피가 b·f가 충분히 크면 완전 무작위와 거의 동일함을 증명한다. 실험에서는 b=16, f=10을 기본값으로 사용했으며, 이는 64‑셀 미니배치당 4 개의 랜덤 디스크 접근과 640 셀의 연속 읽기로 구현된다.

성능 평가에서는 Tahoe‑100M 데이터셋(희소 행렬 300 GB, 밀집 변환 시 1.2 TB)에서 기존 무작위 샘플링이 초당 ~20 셀에 그치는 반면, scDataset는 초당 ~2 500 셀(≈125 배)로 처리량을 끌어올렸다. 또한 AnnLoader, scDataLoader, AnnBatch 등 기존 도구와 비교했을 때 메모리 사용량은 변함없이 낮고, 멀티프로세싱 지원으로 CPU 코어 활용도가 높았다.

다운스트림 작업으로는 세 가지 분류 태스크(세포 타입, 약물 반응, 복합 조건)에서 동일한 모델 아키텍처를 사용했으며, scDataset와 완전 무작위 샘플링 간 정확도 차이는 0.1 % 이하였다. 반면 학습 시간은 무작위 샘플링 대비 30 ~ 40 배 단축되었다.

호환성 측면에서 scDataset는 AnnData, Zarr, HuggingFace Datasets, NumPy 등 다양한 백엔드와 플러그인 형태의 콜백을 제공한다. 데이터 변환 없이 바로 .h5ad 파일을 읽어들일 수 있어 기존 scverse 파이프라인과 무결점 연동이 가능하고, 필요 시 TileDB‑SOMA나 BioNeMo‑SCDL 같은 메모리‑맵 형식에도 적용할 수 있다.

요약하면, scDataset는 “준무작위” 샘플링이라는 새로운 패러다임을 제시함으로써 대규모 단일세포 오믹스 데이터의 딥러닝 학습을 기존 하드웨어 환경에서도 실현 가능하게 만든다.