Token Merging (ToMe)

토큰 병합 (ToMe)은 순전파 중 처리되는 토큰 수를 줄여 Transformer 아키텍처의 성능과 효율성을 최적화하도록 설계된 최첨단 기법입니다. 처음에는 비전 Transformer (ViT) 모델의 처리 속도를 높이기 위해 개발되었으며, ToMe는 추가 학습 없이 네트워크 내 중복 토큰을 체계적으로 식별하고 결합합니다. 셀프 어텐션 메커니즘의 계산 복잡도는 토큰 수에 대해 이차적으로 증가하므로, 유사한 토큰을 병합하면 총 부동 소수점 연산(FLOPs)이 크게 감소하여 실시간 추론 속도를 크게 높일 수 있습니다.

토큰 병합 프로세스 이해하기#

ToMe는 이미지를 개별 토큰으로 나누는 초기 전처리 단계인 토큰화와는 근본적으로 다릅니다. 토큰화가 이산 요소를 생성하는 반면, 토큰 병합은 모델의 순전파 실행 중 동적으로 다운샘플링하는 메커니즘으로 작동합니다.

이 알고리즘은 일반적으로 이분 매칭을 사용하여 토큰 유사도를 평가하며, 어텐션 레이어에서 토큰의 키 사이의 코사인 유사도를 계산하는 경우가 많습니다. 시각적 또는 의미론적 정보가 매우 유사한 토큰은 하나로 융합되며, 이때 특징을 평균내는 방식을 자주 사용합니다. 이를 통해 필수적인 공간 또는 문맥 정보는 보존하면서 불필요한 계산 부담은 줄일 수 있으므로, PyTorch와 같은 프레임워크에서 복잡한 비전 모델을 훨씬 빠르게 처리할 수 있습니다.

토큰 병합의 실제 활용 사례#

토큰 병합은 계산 리소스가 제한된 환경에 대규모 어텐션 기반 아키텍처를 배포하기 위한 핵심 최적화 전략이 되었습니다.

1. 생성형 AI 및 이미지 합성: 널리 사용되는 텍스트-이미지 확산 모델에서 ToMe는 이미지 생성을 가속하는 데 자주 사용됩니다. 배경 또는 세부 정보가 적은 토큰을 병합하면 생성 과정에 필요한 단계 수가 줄어들어 GPU 리소스를 크게 절약하고 생성 모델을 사용하는 최종 사용자의 지연 시간을 줄일 수 있습니다. 확산 프로세스에 대한 자세한 내용은 arXiv의 기초 연구에서 확인할 수 있습니다.

2. 엣지 AI 배포: 세그먼트 애니씽 모델 (SAM)과 같은 대규모 모델을 모바일 장치에 배포하는 작업은 메모리 제약으로 인해 매우 어렵습니다. ToMe는 메모리 사용량을 동적으로 줄여 복잡한 이미지 세그멘테이션 작업을 엣지 하드웨어에서 실행할 수 있도록 지원합니다. 순수한 속도가 중요한 시나리오에서는 더 빠른 엔드투엔드 엣지 추론을 위해 엔지니어들이 Ultralytics YOLO26과 같은 네이티브 최적화 어텐션 없는 아키텍처로 전환하는 경우가 많습니다.

Python 예제: 토큰 유사도 계산#

ToMe를 전체 아키텍처에 통합하려면 어텐션 블록을 수정해야 하지만, 핵심 개념은 유사한 토큰을 찾는 데 있습니다. 다음 PyTorch 코드 조각은 병합 후보가 될 토큰을 식별하기 위해 토큰 집합 간의 코사인 유사도를 계산하는 방법을 보여줍니다.

import torch
import torch.nn.functional as F

# Simulate a batch of 4 image patches (tokens) with 64-dimensional features
tokens = torch.randn(1, 4, 64)

# Normalize the tokens to easily compute cosine similarity via dot product
normalized_tokens = F.normalize(tokens, p=2, dim=-1)

# Compute the similarity matrix between all tokens (1 x 4 x 4)
similarity_matrix = torch.matmul(normalized_tokens, normalized_tokens.transpose(1, 2))

# Tokens with high similarity scores (close to 1.0) off the diagonal
# are prime candidates for Token Merging.
print("Similarity Matrix:", similarity_matrix)

최신 머신 러닝 파이프라인에서는 정확도와 속도의 균형을 신중하게 조정해야 합니다. 사용자 지정 ViT를 최적화하기 위해 토큰 병합을 사용하든 YOLO26의 최첨단 효율성에 의존하든, 복잡한 데이터 워크플로를 관리하는 작업은 Ultralytics Platform을 통해 크게 간소화됩니다. Platform은 자동화된 데이터 어노테이션, 원활한 클라우드 학습, 다양한 엣지 컴퓨팅 하드웨어 환경 전반에 걸친 강력한 모델 배포를 위한 직관적인 생태계를 제공합니다. 컴퓨터 비전 이니셔티브를 확장하는 조직은 이러한 도구를 활용하여 최첨단 모델을 안정적이고 효율적으로 프로덕션에 배포합니다.