안녕하세요 r/LocalLLaMA 여러분! 저희의 첫 번째 VLM인 SupraVL-Nano-900k를 방금 출시했습니다. 이 모델은 약 900k (90만) 개의 파라미터(parameters)를 가지고 있으며, Flickr8k 데이터셋을 통해 처음부터(from scratch) 학습되었습니다. 전체 아키텍처(architecture)는 단일 Jupyter notebook에 들어갈 정도로 가볍습니다. 이것은 상용 서비스용 모델이 아니라, 이미지-텍스트(image-to-text) 모델이 내부적으로 실제로 어떻게 작동하는지 이해하고 싶은 누구에게나 완전히 투명하고 읽기 쉬운 청사진을 제공하기 위한 것입니다.

🤗 SupraVL-Nano-900k

이것은 무엇인가요?
대부분의 VLM은 블랙박스(black boxes)입니다. CLIP 인코더(encoders), 수십억 개의 파라미터를 가진 LLM, 쉽게 읽을 수 없는 융합 계층(fusion layers) 등이 그것입니다. SupraVL-Nano는 이 모든 것을 처음부터 구축합니다: CNN 시각 인코더(visual encoder), GPT-2 스타일의 트랜스포머 디코더(transformer decoder), 데이터셋 자체로 학습된 BPE 토크나이저(tokenizer), 그리고 접두사 결합(prefix concatenation) 융합 전략을 사용합니다. 모든 구성 요소는 처음부터 작성되었으며 문서화되어 있습니다.
목표는 간단합니다: VLM이 어떻게 작동하는지 이해하고 싶다면, 코드를 읽을 수 있어야 한다는 것입니다.

아키텍처(Architecture)

[IMG:1]

구성 요소 상세 정보

• 시각 인코더(Visual encoder): 3× Conv-BN-ReLU + AdaptiveAvgPool(4×4) → 16개의 공간 토큰(spatial tokens)
• 시각 채널(Visual channels): 64-d → 128-d로 투영(projected)
• 디코더(Decoder): GPT-2 스타일, 3개 계층(layers), d=128, 4개 헤드(heads), FF=256
• 어휘집(Vocabulary): Flickr8k 캡션으로 학습된 2048개의 BPE 토큰
• 컨텍스트(Context): 16개의 시각 토큰 + 48개의 텍스트 토큰 = 총 64개 위치
• 파라미터(Parameters): 약 900k
• 융합(Fusion): 접두사 결합(Prefix concatenation) (시각 토큰을 텍스트 시퀀스 앞에 추가)
• 가중치 공유(Weight tying): tok_emb ↔ lm_head (GPT-2 스타일)

4×4 공간 그리드(spatial grid)를 선택한 것은 단일 글로벌 토큰(single global token) 대신 의도적으로 선택한 사항입니다. 이를 통해 디코더가 서로 다른 단어를 생성할 때 이미지의 서로 다른 영역에 주의(attend)를 기울일 수 있으며, 이는 실제 VLM이 작동하는 방식에 더 가깝습니다.

학습 (Training)

설정 값 (Setting Value)

데이터셋 (Dataset): Flickr8k (30k train / 5k val pairs)
에폭 (Epochs): 15
옵티마이저 (Optimizer): AdamW (β₁=0.9, β₂=0.95, wd=0.01)
학습률 (Learning rate): 3e-4 → 코사인 감쇠 (cosine decay) → 3e-5
배치 크기 (Batch size): 64
정밀도 (Precision): 혼합 정밀도 (Mixed (AMP))
하드웨어 (Hardware): Kaggle 2× T4 / Google Colab T4

빠른 시작 (Quick start)
설치 (Install):
pip install torch torchvision pillow huggingface_hub safetensors tokenizers
실행 (Run):
import json, torch, torch.nn as nn, torch.nn.functional as F import torchvision.transforms as T from PIL import Image from huggingface_hub import hf_hub_download from safetensors.torch import load_file from tokenizers import Tokenizer REPO = "SupraLabs/SupraVL-Nano-900k" ckpt_path = hf_hub_download(REPO, "model.safetensors") tok_path = hf_hub_download(REPO, "tokenizer.json") cfg_path = hf_hub_download(REPO, "config.json") with open(cfg_path) as f: cfg = json.load(f) tokenizer = Tokenizer.from_file(tok_path) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 설정 키 (Config keys): D_MODEL, N_HEADS, N_LAYERS, D_FF, VIS_CH, N_VIS, VOCAB_SIZE, MAX_SEQ, IMG_SIZE N_EMBD = cfg["D_MODEL"] # 128 N_HEAD = cfg["N_HEADS"] # 4 N_LAYER = cfg["N_LAYERS"] # 3 D_FF = cfg["D_FF"] # 256 VIS_CH = cfg["VIS_CH"] # 64 VIS_TOKENS = cfg["N_VIS"] # 16 VOCAB_SIZE = cfg["VOCAB_SIZE"] # 2048 MAX_SEQ = cfg["MAX_SEQ"] # 48 IMG_SIZE = cfg["IMG_SIZE"] # 112 TOTAL_POS = VIS_TOKENS + MAX_SEQ # 64 BOS_ID = cfg.get("bos_token_id", 1) EOS_ID = cfg.get("eos_token_id", 2) # --- 모델 정의 (Model definition) --- class CausalSelfAttention(nn.Module): def init(self): super().init() self.qkv = nn.Linear(N_EMBD, 3 * N_EMBD, bias=False) self.proj = nn.Linear(N_EMBD, N_EMBD, bias=False) self.n_head = N_HEAD self.register_buffer( "mask", torch.tril(torch.ones(TOTAL_POS, TOTAL_POS)).view(1, 1, TOTAL_POS, TOTAL_POS) ) def forward(self, x): B, T, C = x.shape nh, hs = self.n_head, C // self.n_head q, k, v = self.qkv(x).split(C, dim=-1) q = q.view(B,T,nh,hs).transpose(1,2) k = k.view(B,T,nh,hs).transpose(1,2) v =

v.view(B,T,nh,hs).transpose(1,2) att = (q @ k.transpose(-2,-1)) * (hs**-0.5) att = att.masked_fill(self.mask[:,:,:T,:T]==0, float("-inf")) att = F.softmax(att, dim=-1)
return self.proj((att @ v).transpose(1,2).contiguous().view(B,T,C))
class MLP(nn.Module):
def init(self):
super().init()
self.fc1 = nn.Linear(N_EMBD, D_FF)
self.fc2 = nn.Linear(D_FF, N_EMBD)
def forward(self, x):
return self.fc2(F.gelu(self.fc1(x))))
class Block(nn.Module):
def init(self):
super().init()
self.ln1 = nn.LayerNorm(N_EMBD)
self.attn = CausalSelfAttention()
self.ln2 = nn.LayerNorm(N_EMBD)
self.mlp = MLP()
def forward(self, x):
x = x + self.attn(self.ln1(x))
x = x + self.mlp(self.ln2(x))
return x
class VisualEncoder(nn.Module):
def init(self):
super().init()
c1,c2,c3 = VIS_CH//4, VIS_CH//2, VIS_CH
self.conv1 = nn.Sequential(nn.Conv2d(3,c1,3,2,1), nn.BatchNorm2d(c1), nn.ReLU(True))
self.conv2 = nn.Sequential(nn.Conv2d(c1,c2,3,2,1), nn.BatchNorm2d(c2), nn.ReLU(True))
self.conv3 = nn.Sequential(nn.Conv2d(c2,c3,3,2,1), nn.BatchNorm2d(c3), nn.ReLU(True))
grid = int(VIS_TOKENS**0.5)
self.pool = nn.AdaptiveAvgPool2d((grid, grid))
self.proj = nn.Linear(c3, N_EMBD)
def forward(self, x):
x = self.conv3(self.conv2(self.conv1(x)))
B,C,H,W = self.pool(x).shape
x = self.pool(x).view(B,C,H*W).transpose(1,2)
return self.proj(x)
class MiniVLM(nn.Module):
def init(self):
super().init()
self.vis_enc = VisualEncoder()
self.tok_emb = nn.Embedding(VOCAB_SIZE, N_EMBD)
self.pos_emb = nn.Embedding(TOTAL_POS, N_EMBD)
self.blocks = nn.ModuleList([Block() for _ in range(N_LAYER)])
self.ln_f = nn.LayerNorm(N_EMBD)
self.lm_head = nn.Linear(N_EMBD, VOCAB_SIZE, bias=False)
def forward(self, img_tokens, tok_ids):
B, T = tok_ids.shape
seq = torch.cat([img_tokens, self.tok_emb(tok_ids)], dim=1)
pos = self.pos_emb(torch.arange(VIS_TOKENS+T, device=tok_ids.device))
x = seq + pos.unsqueeze(0)
for block in self.blocks:
x = block(x)
return self.lm_head(self.ln_f(x))
u/torch.no_grad()

generate_beam(self, img, beam_width=3, max_new=48):
self.eval()
img_tokens = self.vis_enc(img)
beams = [(0.0, [BOS_ID])]
for _ in range(max_new):
candidates = []
for score, seq in beams:
if seq[-1] == EOS_ID:
candidates.append((score, seq)); continue
ids = torch.tensor([seq], dtype=torch.long, device=img.device)
logits = self.forward(img_tokens, ids)
lprobs = F.log_softmax(logits[0, VIS_TOKENS+len(seq)-1], dim=-1)
topk = torch.topk(lprobs, beam_width)
for lp, tok in zip(topk.values.tolist(), topk.indices.tolist()):
candidates.append((score+lp, seq+[tok]))
beams = sorted(candidates, key=lambda x: x[0], reverse=True)[:beam_width]
if all(s[-1]==EOS_ID for _,s in beams):
break
best = [t for t in beams[0][1] if t not in (BOS_ID, EOS_ID)]
return tokenizer.decode(best)

--- 가중치 로드 ---

model = MiniVLM()
model.load_state_dict(load_file(ckpt_path, device=str(device)), strict=False)
model.lm_head.weight = model.tok_emb.weight
model.to(device).eval()

--- 추론 실행 ---

transform = T.Compose([
T.Resize((IMG_SIZE, IMG_SIZE)),
T.ToTensor(),
T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
img = Image.open("your_image.jpg").convert("RGB")
img_t = transform(img).unsqueeze(0).to(device)
print("Caption:", model.generate_beam(img_t, beam_width=3, max_new=48))
Generation strategies

Method Notes

Greedy model.generate_greedy(img) — 빠르고 결정론적 (deterministic)
Top-k sampling model.generate_topk(img, temperature=0.8, top_k=50) — 더 다양한 결과 생성
Beam search model.generate_beam(img, beam_width=3) — 가장 유창하며 권장됨 (recommended)
Limitations (be honest with yourselves)
This is trained on Flickr8k in under an hour on a T4. Expect short generic captions, repetition on out-of-distribution images, nonsense outputs some times and no instruction following whatsoever. It is not competing with LLaVA. It is competing with nothing, it's an educational artifact.

로드맵 (Roadmap)

• CNN을 아주 작은 ViT 패치 인코더 (patch encoder)로 교체
• 접두사 연결 (prefix concatenation) 대신 교차 주의 집중 (Cross-attention) 레이어 사용 (Flamingo 스타일)
• 사전 학습된 동결된 (frozen) CLIP 백본 (backbone) 사용
• 디코더 (decoder)를 6-12개 레이어, d=512+로 확장
• CC3M / LAION-400M 데이터셋으로 학습
• 규모 확장 (Scale up)

Apache 2.0. 가서 코드를 읽어보세요. 그것이 이 프로젝트의 핵심입니다.
제출자: /u/Dangerous_Try3619
[link] [comments]