Flamingo, LLaVA, BLIP-2를 통해 VLM의 연결 구조 정리

VLM (Vision-Language Model)의 구현을 읽을 때, 가장 먼저 확인하고 싶은 것은 model 이름이 아니라 Vision Encoder와 LLM의 인터페이스입니다.

이미지는 보통 다음과 같은 경로로 LLM에 전달됩니다.

pixels
-> Vision Encoder
-> visual features: (B, N, D_v)
...

Connector가 단순한 Linear/MLP라면, visual tokens을 text tokens에 연결할 수 있습니다. Resampler나 Q-Former라면 token 수를 고정 길이로 압축할 수 있습니다. Cross-Attention이라면 language stream과 visual memory를 분리된 상태로 융합할 수 있습니다.

이 글에서는 LLaVA, Flamingo, BLIP-2를 tensor shape와 Attention 경로의 차이점으로부터 정리합니다.

VLM을 세 가지 컴포넌트로 나누기

생성형 VLM은 개념적으로 다음 3부분으로 읽을 수 있습니다.

• Vision Encoder: 이미지를 위치가 지정된 visual features로 변환하는 부분
• Connector: feature dimension, token 수, fusion 방법을 LLM 측에 맞추는 부분
• Language Model: visual condition과 instruction으로부터 answer tokens를 생성하는 부분

주의할 점은 VLM이 하나의 고정된 아키텍처 이름이 아니라는 것입니다. CLIP 같은 Dual Encoder도 넓은 의미에서는 VLM이지만, CLIP 단독의 주 출력은 이미지와 텍스트의 유사도입니다. LLaVA나 Flamingo 같은 생성형 VLM은 시각 정보를 이용해 language model을 조건화하여 자유 형식 텍스트를 생성합니다.

336×336 이미지에서 576개의 visual features 만들기

입력 이미지를 336×336으로, Vision Encoder의 patch size를 14로 가정하겠습니다. patch grid는 24×24이므로, patch feature 수는

batch size는
LLM의 embedding dimension은

여기서 576은 sequence length, 4096은 각 token의 vector width입니다. token 수

이 수치는 메커니즘을 보여주는 가상 예시입니다. 실제 token 수는 Class Token의 처리 방식, Vision Encoder의 레이어, patch merge, tiling, processor 설정에 따라 달라집니다.

Projector의 shape 일치만으로는 이미지를 읽을 수 없다

Linear Projector는 tensor 연산 상 마지막 dimension을
구별해야 합니다. 중요한 것은 다음 2가지입니다.

원래 LLaVA의 feature alignment 단계에서는 CLIP Vision Encoder와 LLM을 동결(freeze)하고, projection matrix만 학습합니다. 원 논문은 이 단계를 동결된 LLM용 호환 visual tokenizer를 배우는 처리로 설명했습니다.

다만, text tokenizer처럼 이미지를 이산 token ID로 변환하는 것은 아닙니다. LLM에 전달되는 것은 연속 embedding입니다.

Projector + concat을 최소 구현하기

다음 코드는 Vision Encoder가 이미 출력한 (B, N, D_v)를 LLM 폭으로 투영하고, (B, T, D_l)의 text embeddings보다 앞에 연결하는 최소 예시입니다. 이미지 processor, Vision Encoder 본체, Position ID, Attention mask, label mask는 생략했습니다.

import logging
import torch
from torch import Tensor, nn
...

실제 model에서는, input_ids

토큰 임베딩 레이어(token embedding layer)로 통과시킨 후, 이미지 placeholder의 위치를 visual embeddings로 대체하는 구현도 있습니다. 이 경우에도 확인해야 할 계약은 동일합니다.

• visual features의 마지막 dimension이 Projector 입력과 일치하는지
• Projector 출력이 LLM token embedding과 같은 dtype・device・width인지
• visual tokens을 배치할 위치와 Attention mask가 대응하는지
• labels 중 loss를 계산하지 않는 위치가 올바르게 mask되었는지

Concat형에서는 visual tokens도 context를 사용한다

visual tokens을 prompt 앞에 배치하는 개념 예시는 다음과 같습니다.

[V1][V2]...[VN][USER][질문...][ASSISTANT][답변...]

deconder-only LLM의 Causal Self-Attention에서는, 답변 위치(answer position)부터 그 이전에 있는 visual tokens, instruction, 과거 answer tokens을 참조할 수 있습니다. 미래의 answer token은 참조할 수 없습니다.

이미지

training 시에는 answer 위치만 label로 사용하는 레시피가 있습니다. 반면, 어느 prompt 위치를 loss에서 제외할지는 model 구현과 data format에 따라 다르므로, labels == -100와 같은 mask를 실제 데이터(real data)에서 확인해야 합니다.

concat형은 기존 Causal LLM에 통합하기 쉽다는 장점이 있지만, visual tokens 역시 LLM context를 소비합니다. 이미지가 576 tokens이고 텍스트가 32 tokens이라면 prefill sequence는 608입니다. 여러 이미지를 사용하는 경우 단순화하면 다음과 같이 증가합니다.

Cross-Attention형은 visual memory를 분리한다

Cross-Attention에서는 language hidden state에서 Query를 만들고, visual features를 Key・Value로 참조합니다.

concat형에서는 visual tokens과 text tokens이 같은 Self-Attention sequence에 들어갑니다. Cross-Attention형에서는 language stream의 길이와 visual memory의 길이를 분리할 수 있습니다.

어느 쪽이 항상 우월하다고 할 수 있는 관계는 아닙니다. 구현의 단순성, 여러 이미지 처리 여부, visual token budget, 기존 LLM 재사용 방법 등이 선택 기준입니다.

Flamingo는 64 tokens으로 Resample하여 LM block 사이에 삽입한다

Flamingo는 동결된 Vision Encoder의 가변 길이 spatial・temporal features를 Perceiver Resampler에 입력합니다. learned latent queries가 features에 attention하고, 이미지/비디오마다 고정된 64개의 visual tokens을 출력합니다.

다음으로, 동결된 LM block 사이에 Gated Cross-Attention 레이어를 삽입합니다. text hidden states가 Q이고, 64 visual tokens이 K/V입니다.

추가 모듈의 출력은 다음 형태로 residual stream에 더해집니다.

64는 Flamingo의 설계 값이며, VLM 일반의 token 수는 아닙니다. 고정 길이화는 Cross-Attention 비용을 제어하는 한편, 세밀한 문자나 다수의 object를 제한된 표현으로 압축합니다.

BLIP-2는 32 learned queries를 정보 병목(information bottleneck)으로 사용한다

BLIP-2는 동결된 Image Encoder와 동결된 LLM 사이에 Q-Former (Querying Transformer)를 배치합니다.

원 논문의 예시에서는, 32 learned queries가 ViT-L/14의 257×1024

features에 Cross-Attention을 수행하여, 32×768의 query outputs를 만듭니다.

Q-Former는 2단계로 pre-training됩니다.

• Vision-Language Representation Learning: ITC, ITG, ITM을 사용하여 text와 관련된 visual representation을 query에 모읍니다.
• Vision-to-Language Generative Learning: query outputs를 LLM의 차원(dimension)으로 투영(project)하고, soft visual prompts로 고정된 LLM에 연결합니다.

입력 resolution에 따라 ViT feature 수가 달라져도 Q-Former 출력 수는 32로 고정할 수 있습니다. 다만, 고정 query 수는 정보 병목 현상(information bottleneck)입니다. document OCR처럼 국소적인 디테일이 많은 task에서는 query 수, input resolution, tiling, training data를 함께 평가해야 합니다.

Fusion 위치에서 아키텍처 비교하기

대표 모델들을 '어디서 vision과 language를 융합할지'에 따라 정리합니다.

Checkpoint는 4개 포인트를 한 세트로 취급하기

VLM을 로컬에서 불러올 때는 LLM weights뿐만 아니라 다음 조합을 갖추어야 합니다.

• Vision Encoder checkpoint
• Image processor / preprocessor
• Connector checkpoint
• Language model checkpoint와 chat template

dimension이 일치하는 다른 Projector로 교체해도 semantic alignment는 일치하지 않습니다. 또한 processor가 다르면, resize, crop, normalization, tile 순서가 training 시점과 달라집니다.

구현 조사에서는 다음 사항을 확인합니다.

• input resolution과 dynamic tiling의 유무
• patch size, patch merge, Class Token 제외 여부
• visual token 수가 고정인지 입력에 의존하는지
• image placeholder를 몇 개의 embedding으로 전개(expand)할지
• Position ID를 visual tokens에 어떻게 할당할지
• Attention mask가 여러 이미지를 어떻게 구별할지
• Vision Encoder, Connector, LLM의 freeze 정책
• dtype, quantization, device placement가 컴포넌트 간 일관적인지
• context 상한선에 visual tokens을 포함할지

특히 quantization에서는, LLM만 4bit화하더라도 Vision Encoder와 Connector의 activation 및 memory는 남아 있습니다. latency를 측정할 때는 image preprocessing, vision encode, prefill, decode를 분리하면 병목 현상을 찾기 쉽습니다.

유창함(Fluency)과 visual grounding을 별도로 평가하기

VLM의 출력은 visual evidence와 language prior 양쪽 모두에 의해 결정됩니다. 작은 글자가 뭉개졌거나, object가 patch보다 작거나, Connector가 디테일을 압축한 경우 등에는 LLM이 그럴듯한 단어를 보충할 수 있습니다.

평가를 하나의 chat 품질 점수로 통합하지 말고, 최소한 다음 슬라이스(slice)로 나누어 진행해야 합니다.

• object existence: 존재하지 않는 객체를 언급하는지
• counting: 개수를 유지할 수 있는지
• spatial relation: 좌우, 상하, 포함 관계를 혼동하지 않는지
• OCR: 소문자, 회전 문자, 표 구조를 오독하지 않는지
• attributes: 색상, 재질, 상태를 혼동하지 않는지
• abstention: 읽을 수 없을 때 불확실성을 나타낼 수 있는지

자연스러운 설명문을 생성할 수는 있지만, 픽셀 레벨의 위치나 수량이 교정되었다고는 할 수 없습니다. detection, segmentation, 엄밀한 measurement에서는 전용 head 또는 tool과 결합하여 태스크 고유의 ground truth로 검증합니다.

요약

VLM의 구현은 Vision Encoder・Connector・Language Model의 인터페이스에서 읽으면 정리할 수 있습니다.

(B, N, D_v)
을 Projector를 통해 (B, N, D_l)
로 변환하더라도 시맨틱 얼라인먼트에는 트레이닝이 필요합니다. - concat 타입은 visual tokens를 LLM context에 넣고 Causal Self-Attention으로 답변에서 참조합니다.

• Cross-Attention 타입은 text를 Q(쿼리), vision을 K/V(키/값)로 하여 스트림을 분리합니다.
• Flamingo는 Perceiver Resampler를 사용하여 64 토큰으로 고정하고, Gated Cross-Attention을 LM 블록 사이에 삽입합니다.
• BLIP-2는 32개의 학습된 쿼리를 가진 Q-Former를 정보 병목(information bottleneck)으로 사용합니다.
• 고해상도화에서는 detail뿐만 아니라 token 수, prefill, KV Cache, 압축 손실을 측정합니다.

모델 이름보다 먼저 visual feature의 shape, fusion 위치, token budget, freeze 대상을 확인해야 합니다. 이 4가지 점이 파악되면 VLM family 간의 차이점을 코드와 config에서 추적하기 쉬워집니다.

VLM의 전체적인 모습, 학습 단계, 도입 시 설계 판단을 포함한 자세한 해설은 개인 블로그 버전의 'VLM이란?' 이미지를 이해하는 대규모 언어 모델에 정리되어 있습니다.

참고 문헌

토론