nanoVLM: 비전 언어 모델 훈련을 위한 미니멀리스트 PyTorch 툴킷

Hugging Face는 nanoVLM을 소개했습니다. 순수 PyTorch로 작성된 미니멀리스트 툴킷으로, Vision Language Models (VLM) 훈련 과정을 쉽게 이해하도록 설계되었습니다. Andrej Karpathy의 nanoGPT에서 영감을 받은 nanoVLM은 읽기 쉬운 경량 코드베이스를 제공하여 사용자가 무료 티어 Google Colab 노트북에서 VLM을 훈련할 수 있게 합니다.

Vision Language Model (VLM) 기본 개념

Vision Language Model은 이미지와 텍스트 입력을 모두 처리하여 텍스트 출력을 생성할 수 있는 멀티모달 아키텍처입니다. VLM은 이미지 캡션 생성, 객체 탐지, 의미론적 분할 등 다양한 작업에 활용될 수 있지만, nanoVLM은 주된 훈련 목표로 **Visual Question Answering (VQA)**에 특별히 초점을 맞춥니다.

기술 아키텍처

nanoVLM은 두 개의 사전 훈련된 백본을 프로젝션 레이어를 통해 정렬하는 모듈식 아키텍처를 활용합니다:

  • Vision Backbone: Google의 SigLIP (google/siglip-base-patch16-224) 비전 인코더를 사용합니다.
  • Language Backbone: Llama 3 아키텍처를 따르며, 특히 HuggingFaceTB/SmolLM2-135M을 활용합니다.
  • Modality Projection Module: 이 모듈은 비전과 텍스트 모달리티를 정렬합니다. 비전 백본의 이미지 임베딩을 언어 모델의 임베딩 레이어와 호환되는 임베딩으로 변환합니다. 이 과정은 pixel shuffle 연산을 포함하는데—이미지 토큰 수를 줄여 계산 비용을 낮추고 훈련 속도를 높입니다—그 후 선형 레이어가 이어집니다.

훈련 워크플로우 및 구현

train.py를 통해 훈련 프로세스를 관리하며, 데이터셋 로드, 모델 초기화 및 최적화를 수행합니다.

데이터 파이프라인

get_dataloaders 함수는 Hugging Face load_dataset API를 활용하여 데이터셋을 로드하고, 셔플하며, 분할합니다. 맞춤형 데이터셋 (VQADataset, MMStarDataset) 및 콜레이터 (VQACollator, MMStarCollator)를 사용해 데이터를 준비합니다.

최적화 전략

사전 훈련된 백본과 새로 초기화된 프로젝터의 훈련 균형을 맞추기 위해 nanoVLM은 이중 학습률(LR) 전략을 사용합니다:

  • Higher LR: Modality Projector (MP)에 적용되어 빠른 학습을 촉진합니다.
  • Lower LR: 인코더/디코더 스택에 적용되어 백본에 이미 존재하는 지식을 보존합니다.

훈련 루프 및 모니터링

훈련은 혼합 정밀도를 위해 torch.autocast를 사용하고, 코사인 학습률 스케줄에 선형 워밍업을 적용합니다. 성능은 토큰 처리량 (tokens/sec)으로 모니터링되며, 활성화된 경우 Weights & Biases (wandb)를 통해 배치 손실, 검증 손실 및 정확도를 추적합니다.

추론 및 사전 훈련 모델

Hugging Face는 Hub에 공개된 사전 훈련 nanoVLM 모델 (nanoVLM-222M)을 제공했습니다. 이 모델은 cauldron 데이터셋에서 170만 개 샘플을 사용하여 단일 H100 GPU에서 약 6시간 동안 훈련되었습니다.

사용자는 generate.py 스크립트를 사용해 추론을 실행할 수 있으며, 다음과 같은 논리 흐름을 따릅니다:

  1. Initialization: 모델, 토크나이저 및 이미지 프로세서를 로드합니다.
  2. Processing: 텍스트 프롬프트를 토크나이즈하고 이미지를 텐서로 처리합니다.
  3. Generation: model.generate를 실행하여 텍스트 출력을 생성합니다.
  4. Decoding: batch_decode를 사용해 생성된 토큰을 인간이 읽을 수 있는 텍스트로 변환합니다.

시작하기

훈련을 시작하려면, 사용자는 저장소를 복제하고 훈련 스크립트를 실행하면 됩니다:

# Clone the repo
git clone https://github.com/huggingface/nanoVLM.git

# Execute the training script
python train.py

Sources