Token Merging (ToMe)

토큰 병합(ToMe)은 네트워크 내부의 유사한 토큰을 결합하여 각 블록이 처리하는 토큰 수를 줄이는 방식으로 트랜스포머, 특히 Vision Transformer(ViT)의 속도를 높입니다. 이미지에는 거의 동일하게 보이는 하늘이나 도로 패치처럼 중복 정보가 많습니다. 이런 토큰을 평균하여 결합하면 정보를 유지하면서 연산량을 줄일 수 있습니다. ToMe는 재학습 없이 이미 학습된 모델에 적용할 수 있으며, 학습 중에도 사용할 수 있습니다. Georgia Tech와 Meta AI의 Daniel Bolya와 동료들은 ICLR 2023 구두 발표 논문인 Token Merging: Your ViT But Faster에서 이 방법을 소개했습니다. 별도 설정 없이 적용했을 때 이미지 정확도 하락을 0.2%~0.3%로 제한하면서 512픽셀 ViT-L과 518픽셀 ViT-H의 처리량을 두 배로 높였으며, 비슷한 정확도 하락으로 ViT-L 비디오 모델의 속도를 2.2배 높였습니다.

토큰 병합 작동 방식#

ViT는 이미지를 패치로 나누어 패치마다 하나의 토큰을 할당합니다. 각 블록은 토큰 수의 제곱에 따라 비용이 늘어나는 셀프 어텐션 다음에 비용이 선형으로 늘어나는 MLP를 실행합니다. ToMe는 각 블록의 어텐션과 MLP 사이에 병합 단계를 삽입하고 매번 고정된 수의 토큰 r개를 제거합니다. 따라서 r = 8인 24블록 모델은 시작 시보다 토큰이 192개 적은 상태로 종료됩니다. 병합에는 빠르고 반복적 클러스터링이 필요하지 않아 선택된 이분 그래프 소프트 매칭을 사용합니다.

1. 분할: 토큰을 번갈아 두 집합 A와 B로 나눕니다.
2. 매칭: 어텐션 헤드 전체에서 평균한 어텐션 키의 코사인 유사도로 측정하여, A의 각 토큰을 B에서 가장 유사한 토큰과 연결합니다.
3. 최적 엣지 선택: 유사도가 가장 높은 r개의 쌍을 선택합니다.
4. 병합: 각 토큰이 나타내는 원래 패치 수를 가중치로 사용해 선택된 각 쌍을 하나의 토큰으로 평균한 뒤, 두 집합을 다시 이어 붙입니다.

병합된 토큰은 여러 패치를 나타내므로 ToMe는 비례 어텐션이라고 하는 방식으로 각 토큰 크기의 로그를 어텐션 점수에 더합니다. 이를 통해 여러 패치로 만들어진 토큰이 그에 상응하는 어텐션을 받습니다. 전체 절차는 미분 가능하므로 ToMe를 적용해 모델을 학습하고 추가 속도 향상을 얻을 수 있습니다.

이 PyTorch 예제는 448 x 448 버스 이미지에서 Hugging Face Transformers의 Meta DINOv2-small ViT에 이분 그래프 소프트 매칭을 적용합니다. 12개 블록 각각 뒤에서 r개의 토큰을 병합하고, 그 결과 이미지 임베딩을 병합하지 않은 모델의 임베딩과 비교합니다.

from urllib.request import urlopen

import torch
import torch.nn.functional as F
from PIL import Image
from transformers import AutoImageProcessor, AutoModel

processor = AutoImageProcessor.from_pretrained("facebook/dinov2-small")
model = AutoModel.from_pretrained("facebook/dinov2-small")
image = Image.open(urlopen("https://ultralytics.com/images/bus.jpg")).resize((448, 448))
pixels = processor(images=image, do_resize=False, do_center_crop=False, return_tensors="pt").pixel_values


def merge(x, r):
    """Bipartite soft matching: merge the r most similar pairs between alternating token sets A and B."""
    cls, a, b = x[:, :1], x[:, 1::2], x[:, 2::2]
    similarity = F.normalize(a[0], dim=-1) @ F.normalize(b[0], dim=-1).T
    best, match = similarity.max(dim=-1)  # each A token's most similar B token
    order = best.argsort(descending=True)
    merged, kept = order[:r], order[r:]
    index = match[merged][None, :, None].expand(1, r, x.shape[-1])
    b = b.scatter_reduce(1, index, a[:, merged], reduce="mean")  # average merged A tokens into their B matches
    return torch.cat([cls, a[:, kept], b], dim=1)


with torch.no_grad():
    for r in (0, 32, 64):
        x, counts = model.embeddings(pixels), []
        for layer in model.encoder.layer:
            counts.append(x.shape[1])
            x = merge(layer(x), r)  # merge r tokens after every transformer block
        embedding = model.layernorm(x)[:, 0]  # CLS token: the image embedding
        reference = embedding if r == 0 else reference
        similarity = F.cosine_similarity(embedding, reference).item()
        print(f"r={r}: mean tokens per block {sum(counts) / 12:.0f}, CLS similarity to unmerged {similarity:.3f}")

병합하지 않은 모델은 모든 블록에서 1025개의 토큰을 처리합니다. 이는 1,024개 패치와 CLS 토큰을 합한 수입니다. 블록마다 토큰 32개를 병합하면 평균 토큰 수가 849로 감소하고, 이미지 임베딩은 원본과 0.983의 코사인 유사도를 유지합니다. 블록마다 64개를 병합하면 유사도가 0.887인 상태에서 토큰 수가 673로 감소합니다. 이를 통해 r이 연산량과 충실도 사이의 절충을 어떻게 조정하는지 확인할 수 있습니다.

토큰 병합의 활용 분야#

Token Merging for Fast Stable Diffusion은 Stable Diffusion 디노이징 네트워크의 트랜스포머 블록에 같은 아이디어를 적용합니다. 추가 학습 없이 토큰을 최대 60% 제거하여 이미지를 최대 두 배 빠르게 생성하고 메모리 사용량을 최대 5.6배 줄일 수 있습니다. 이미지 이외에도 ToMe 논문은 병합을 적용해 학습한 오디오 ViT-B의 처리량을 두 배로 높이면서 mAP를 0.4% 낮췄고, 비디오의 마스크드 오토인코더 파인튜닝 속도를 최대 두 배 높였습니다. 연구자들은 비전-언어 모델의 시각 토큰에도 유사한 병합 및 프루닝을 적용합니다. LLaVA-1.5는 이미지마다 언어 모델의 컨텍스트에서 576개 토큰을 사용합니다.

Ultralytics 패키지에서 Meta의 Segment Anything Model(SAM)은 각 이미지를 1024 x 1024 픽셀로 인코딩하여 패치 토큰 4,096개를 생성합니다. MobileSAM 비교에서 ViT-H 이미지 인코더는 GPU에서 이미지당 SAM에 필요한 456 ms 중 452 ms를 차지합니다. 토큰 감소 방법이 겨냥하는 연산량이 바로 이 부분입니다.

토큰 병합과 관련 개념 비교#

• 토큰 프루닝: DynamicViT 및 EViT와 같은 방법은 중요하지 않다고 판단한 토큰을 버립니다. 이 방식은 토큰의 정보를 잃으며 대개 학습이 필요합니다. ToMe는 토큰을 버리는 대신 병합하여 정보를 평균에 반영하며, 토큰 병합 없이 학습된 모델에도 적용할 수 있습니다.
• 토큰화: 네트워크 실행 전에 패치나 텍스트에서 초기 토큰을 만듭니다. ToMe는 네트워크 내부에서 토큰 수를 줄입니다.
• Flash attention 및 희소 어텐션: 토큰 수를 유지하면서 어텐션의 비용을 낮추거나 속도를 높입니다. ToMe는 토큰 수를 줄여 MLP 레이어의 속도도 높이며, 이 방법들과 함께 사용할 수 있습니다.
• Mixture of Depths(MoD): 학습된 라우터가 언어 모델에서 어떤 토큰이 블록을 건너뛸지 결정하도록 합니다. ToMe는 고정된 유사도 규칙으로 중복 토큰을 영구적으로 병합합니다.
• 입력 해상도 낮추기: 이미지 전체에서 토큰 수를 균일하게 줄입니다. ToMe는 배경과 같이 중복이 많은 영역에서 주로 토큰을 병합하며, 세부 영역은 전체 해상도로 유지합니다.