nanoVLM: The simplest repository to train your VLM in pure PyTorch

🖼️ 인포그래픽

🖼️ 4컷 인포그래픽

💡 한 줄 요약

nanoVLM은 SigLIP 비전 인코더와 Llama 3 계열 언어 구조를 간결한 순수 PyTorch 코드로 연결해, 시각 질의응답용 VLM의 구성부터 학습·평가·추론까지 직접 익힐 수 있게 만든 경량 툴킷이다.

📌 핵심 요약

• nanoVLM은 nanoGPT의 단순성과 가독성에서 영감을 받아, 초보자도 비전 언어 모델의 내부 구조를 살펴보고 무료 Colab 환경에서 학습을 시작할 수 있도록 설계됐다.
• 모델은 사전 학습된 google/siglip-base-patch16-224 비전 백본과 HuggingFaceTB/SmolLM2-135M 언어 백본을 사용하며, 픽셀 셔플과 선형 계층으로 구성된 모달리티 프로젝터가 두 임베딩 공간을 연결한다.
• train.py는 데이터셋 로딩과 전처리, 모델 초기화, 두 종류의 학습률을 사용하는 최적화, 혼합 정밀도 학습, 검증, 체크포인트 저장 및 로깅을 하나의 흐름으로 제공한다.
• 학습 중에는 선형 워밍업이 포함된 코사인 학습률 일정과 처리량 로깅을 사용하고, 설정된 간격마다 검증 데이터와 MMStar 테스트 데이터에서 평가해 정확도가 향상되면 모델을 저장한다.
• 공개된 사전 학습 모델은 약 170만 개의 Cauldron 샘플로 단일 H100 GPU에서 약 6시간 학습됐으며, generate.py 또는 Hugging Face Space를 통해 이미지와 질문을 입력하고 생성된 답변을 확인할 수 있다.

🧩 주요 포인트

1. nanoVLM은 nanoGPT의 단순성과 가독성에서 영감을 받아, 초보자도 비전 언어 모델의 내부 구조를 살펴보고 무료 Colab 환경에서 학습을 시작할 수 있도록 설계됐다.
2. 모델은 사전 학습된 google/siglip-base-patch16-224 비전 백본과 HuggingFaceTB/SmolLM2-135M 언어 백본을 사용하며, 픽셀 셔플과 선형 계층으로 구성된 모달리티 프로젝터가 두 임베딩 공간을 연결한다.
3. train.py는 데이터셋 로딩과 전처리, 모델 초기화, 두 종류의 학습률을 사용하는 최적화, 혼합 정밀도 학습, 검증, 체크포인트 저장 및 로깅을 하나의 흐름으로 제공한다.
4. 학습 중에는 선형 워밍업이 포함된 코사인 학습률 일정과 처리량 로깅을 사용하고, 설정된 간격마다 검증 데이터와 MMStar 테스트 데이터에서 평가해 정확도가 향상되면 모델을 저장한다.
5. 공개된 사전 학습 모델은 약 170만 개의 Cauldron 샘플로 단일 H100 GPU에서 약 6시간 학습됐으며, generate.py 또는 Hugging Face Space를 통해 이미지와 질문을 입력하고 생성된 답변을 확인할 수 있다.

🧠 상세 정리

1. nanoVLM의 목표와 설계 철학

nanoVLM은 순수 PyTorch로 자신만의 비전 언어 모델을 학습하는 과정을 가능한 한 단순하게 보여주기 위해 만든 경량 툴킷이다. 프로젝트는 Andrej Karpathy의 nanoGPT에서 영감을 받아, 복잡한 프레임워크나 방대한 추상화 대신 작고 읽기 쉬운 코드로 핵심 원리를 드러내는 방식을 택했다. 따라서 완성된 대형 모델의 성능만 사용하는 사람보다 이미지와 텍스트가 모델 내부에서 어떻게 결합되고 학습되는지 직접 살펴보려는 초보자와 개발자에게 초점을 맞춘다. 저장소를 복제한 뒤 python train.py를 실행하는 간단한 진입점을 제공하며, 로컬 환경을 따로 준비하지 않아도 무료 등급 Colab 노트북에서 학습을 시작할 수 있도록 했다. 글은 이러한 목표에 맞춰 VLM의 개념, 저장소 구성, 모델 아키텍처, 학습 파이프라인, 사전 학습 모델의 추론 방법을 차례로 설명한다.

2. 비전 언어 모델의 역할과 nanoVLM의 범위

비전 언어 모델은 이미지와 텍스트라는 두 종류의 입력을 처리하고, 그 내용을 바탕으로 텍스트를 생성하는 멀티모달 모델이다. 이 구조를 이용하면 이미지 설명 생성, 객체 탐지, 의미론적 분할, 이미지에 관한 질문 응답처럼 시각 정보와 언어 생성을 결합한 여러 작업을 수행할 수 있다. 원문은 침대 위의 고양이를 문장으로 묘사하거나 고양이 수를 답하는 사례를 통해, 같은 이미지라도 질문과 학습 목표에 따라 서로 다른 형태의 출력을 만들 수 있음을 보여준다. 다만 nanoVLM이 실제 학습 목표로 삼는 작업은 이 전체 범위가 아니라 이미지의 내용을 근거로 질문에 답하는 시각 질의응답이다. 따라서 저장소는 범용 멀티모달 시스템의 모든 기능을 구현하기보다, 이미지 표현과 텍스트 표현을 연결해 자연어 답변을 생성하는 핵심 학습 과정을 명확히 드러내는 데 집중한다.

3. 작고 명확하게 분리된 저장소 구조

저장소는 데이터 처리, 모델 구현, 학습, 추론이라는 주요 책임이 파일과 디렉터리별로 분리되어 있어 전체 흐름을 따라가기 쉽게 구성된다. data 디렉터리에는 데이터셋 래퍼를 정의하는 datasets.py, 배치 구성을 담당하는 collators.py, 입력 전처리를 다루는 processors.py가 들어 있다. models 디렉터리에는 설정, 언어 모델, 비전 트랜스포머, 모달리티 프로젝터, 통합 비전 언어 모델 및 공용 유틸리티 구현이 각각 별도 파일로 배치된다. 저장소 최상위의 train.py는 이 구성 요소를 연결해 전체 학습을 수행하고, generate.py는 학습되거나 공개된 모델을 불러와 이미지와 질문으로 답변을 생성한다. 저자들은 각 파일을 가볍고 문서화된 상태로 유지했으며, models 아래의 개별 구현을 직접 읽어보는 것이 모델의 세부 동작을 이해하는 핵심 방법이라고 권한다.

4. 비전·언어 백본과 모달리티 정렬 구조

nanoVLM의 비전 백본은 표준 비전 트랜스포머 계열인 Google SigLIP 인코더를 따르고, 언어 백본은 Llama 3 아키텍처를 따른다. 기본 학습 구성에서는 비전 쪽에 google/siglip-base-patch16-224의 사전 학습 가중치를 사용하고, 언어 쪽에는 HuggingFaceTB/SmolLM2-135M을 사용한다. 비전 백본이 만든 이미지 임베딩은 언어 모델의 텍스트 임베딩과 바로 결합할 수 없으므로, 모달리티 프로젝터가 이미지 표현을 언어 임베딩과 호환되는 공간으로 변환한다. 이 프로젝터는 픽셀 셔플 연산과 선형 계층으로 구성되며, 변환된 이미지 임베딩은 텍스트 임베딩과 이어 붙여진 뒤 언어 디코더에 입력된다. 특히 픽셀 셔플은 이미지 토큰 수를 줄여 입력 길이에 민감한 트랜스포머 언어 디코더의 계산량을 낮추고 학습 속도를 높이며, 필요하면 다른 SigLIP·SigLIP 2 또는 SmolLM2 변형으로 백본을 교체할 수도 있다.

5. 설정, 데이터 로딩과 모델 초기화

train.py는 데이터 로딩과 전처리부터 모델 초기화, 최적화, 평가와 로깅까지 학습 파이프라인 전체를 담당하는 단일 진입점이다. 먼저 models/config.py에서 학습률과 체크포인트 경로 같은 실행 조건을 담은 TrainConfig와, 은닉 차원 및 어텐션 헤드 수처럼 모델 구조를 정의하는 VLMConfig를 불러온다. 데이터 파이프라인의 중심인 get_dataloaders 함수는 Hugging Face의 load_dataset API로 데이터를 읽고, 여러 데이터셋이 지정되면 이를 결합하고 섞은 뒤 인덱스를 이용해 학습용과 검증용으로 나눈다. 이후 데이터를 VQADataset 또는 MMStarDataset으로 감싸고, VQACollator와 MMStarCollator를 이용해 모델에 입력할 배치를 구성하며, data_cutoff_idx로 작은 부분집합만 선택해 디버깅할 수도 있다. 모델은 VisionLanguageModel 클래스로 생성하고, 기존 체크포인트를 이어서 사용할 때는 from_pretrained로 불러오며, 새로 시작할 때는 선택적으로 사전 학습된 비전 및 언어 백본을 탑재한 모델을 초기화한다.

6. 두 학습률을 사용하는 최적화와 평가

모달리티 프로젝터는 새로 초기화되지만 비전 인코더와 언어 디코더에는 이미 사전 학습된 가중치가 있으므로, nanoVLM은 매개변수를 두 그룹으로 나눠 서로 다른 학습률을 적용한다. 프로젝터에는 더 높은 학습률을 사용해 이미지 표현과 텍스트 표현의 연결을 빠르게 학습시키고, 인코더와 디코더에는 더 낮은 학습률을 사용해 기존 백본의 지식을 과도하게 바꾸지 않도록 한다. 학습 루프는 torch.autocast 기반의 혼합 정밀도를 사용하며, 선형 워밍업과 코사인 감소를 결합한 학습률 일정을 get_lr 함수로 계산한다. 각 배치에서는 손실뿐 아니라 초당 처리 토큰 수를 기록해 학습 성능을 관찰하고, 기본적으로 250스텝마다 검증 데이터와 MMStar 테스트 데이터에서 모델을 평가한다. 평가 정확도가 이전보다 좋아지면 체크포인트를 저장하며, log_wandb가 활성화되면 배치 손실, 검증 손실, 정확도, 토큰 처리량을 Weights & Biases에 전송하고 실행 조건과 날짜를 조합해 실행 이름도 자동 생성한다.

7. 모델 저장·공유와 추론 흐름

학습이 끝난 모델은 save_pretrained를 통해 지정한 경로에 저장할 수 있고, push_to_hub를 호출하면 다른 사용자가 찾고 시험할 수 있도록 Hugging Face Hub에 게시할 수 있다. 추론은 generate.py로 수행하며, 기본 실행에서는 assets/image.png를 대상으로 ‘이것은 무엇인가?’라는 질문을 전달해 생성 결과를 출력한다. 사용자 이미지와 질문을 사용하려면 명령줄의 image와 prompt 인자를 지정하면 되고, 스크립트는 먼저 from_pretrained로 모델을 불러온 뒤 평가 모드로 전환한다. 이어 모델 설정에 맞는 토크나이저와 이미지 프로세서를 초기화하고, 질문과 답변 형식의 텍스트를 토큰화하는 동시에 이미지를 RGB로 읽어 배치 차원을 가진 텐서로 변환한다. 마지막으로 model.generate가 이미지 텐서와 질문 토큰을 바탕으로 새로운 토큰을 생성하고, batch_decode가 이를 사람이 읽을 수 있는 답변으로 복원하며, 여러 번 생성하도록 설정하면 각 결과를 차례로 출력한다.

8. 공개 모델의 학습 결과와 프로젝트의 의미

저자들은 nanoVLM 툴킷을 이용해 google/siglip-base-patch16-224와 HuggingFaceTB/SmolLM2-135M을 백본으로 사용하는 모델을 실제로 학습해 Hub에 공개했다. 이 모델은 Cauldron의 약 170만 개 샘플을 사용해 단일 H100 GPU에서 약 6시간 동안 학습됐으며, 예시에서는 침대 위 두 고양이를 묘사하거나 사진 속 여성이 요가를 하고 있다고 답한다. 사용자는 generate.py를 실행하는 것 외에도 제공된 Hugging Face Space의 사용자 인터페이스에서 모델과 직접 상호작용할 수 있다. 원문은 이 모델이 최고 수준의 모델과 성능을 경쟁하기 위한 것이 아니라, VLM의 구성 요소와 학습 과정을 이해하기 쉽게 풀어내기 위한 결과물이라고 분명히 밝힌다. nanoVLM의 핵심 가치는 가벼운 코드베이스 안에서 멀티모달 입력의 정렬, 차등 학습률을 적용한 최적화, 평가와 체크포인트 저장, 실제 텍스트 생성까지 하나의 재현 가능한 흐름으로 확인할 수 있다는 데 있다.

🧾 핵심 주장 / 시사점

• nanoVLM은 비전 인코더, 언어 디코더, 모달리티 프로젝터를 별도 파일로 분리하고 전체 결합 과정을 짧은 학습·추론 스크립트로 연결해, VLM을 완성된 제품이 아니라 추적 가능한 구성 요소의 집합으로 학습할 수 있게 한다.
• 새로 초기화된 모달리티 프로젝터에는 높은 학습률을, 사전 학습된 비전·언어 백본에는 낮은 학습률을 적용하는 방식은 새 연결 계층의 빠른 적응과 기존 표현 지식의 보존을 동시에 고려한 핵심 학습 전략이다.
• 공개 모델의 목적과 한계를 최고 성능 경쟁이 아닌 구조와 학습 과정의 이해로 명시함으로써, 프로젝트의 성공 기준을 모델 규모나 벤치마크 순위보다 코드 가독성, 실행 용이성, 재현 가능한 학습 흐름에 두고 있다.

✅ 액션 아이템

• nanoVLM 파이프라인을 기준으로 train.py에서 데이터셋 로딩, 전처리, 초기화, 혼합 정밀도 학습, 검증, 체크포인트 저장 흐름을 연결해 재현한다.
• google/siglip-base-patch16-224와 HuggingFaceTB/SmolLM2-135M 결합 구조를 기준으로 모달리티 프로젝터(픽셀 셔플+선형 계층) 설정을 고정하고 두 종류 학습률 최적화의 동작을 점검한다.
• 약 170만 Cauldron 샘플로 공개된 사전학습 모델 조건을 기준으로 단일 H100 6시간 구간, MMStar 평가, generate.py/Space 추론 결과를 동일 실험세트로 비교한다.

❓ 열린 질문

• Colab 무료 환경에서 단일 H100 6시간 학습 재현이 실제로 어느 수준까지 가능한지, 어떤 병목이 우선 드러나는가?
• MMStar 평가 주기에서 정확도 상승 시점을 체크포인트 저장 기준으로 잡는 최적값은 무엇이며 어떻게 판단할 것인가?
• 픽셀 셔플+선형 계층 모달리티 프로젝터가 이미지-질문 간 임베딩 정합에 미치는 한계는 어떤 기준으로 측정할 것인가?