Token Merging (ToMe)

토큰 병합(ToMe)은 순전파(forward pass) 중에 처리되는 토큰 수를 줄여 Transformer 아키텍처의 성능과 효율성을 최적화하도록 설계된 최첨단 기술입니다. 원래 Vision Transformer (ViT) 모델을 가속하기 위해 개발된 ToMe는 추가적인 학습 없이도 네트워크 내의 중복 토큰을 체계적으로 식별하고 결합하여 작동합니다. self-attention 메커니즘의 계산 복잡도는 토큰 수에 따라 이차적으로 증가하므로, 유사한 토큰을 병합하면 전체 부동 소수점 연산(FLOPs)이 크게 줄어들어 훨씬 빠른 real-time inference가 가능해집니다.

Link to this section토큰 병합 프로세스 이해#

ToMe is fundamentally different from tokenization, which is the initial preprocessing step of breaking down an image or text into individual tokens. While tokenization creates the discrete elements, Token Merging acts as a dynamic downsampling mechanism during the model's forward execution.

이 알고리즘은 일반적으로 이분 매칭(bipartite matching)을 사용하여 토큰 유사도를 평가하며, 종종 어텐션 레이어 내 토큰의 키(key) 간 코사인 유사도를 계산합니다. 시각적 또는 의미론적 정보가 매우 유사한 토큰들은 기능을 평균화하는 등의 방식으로 병합됩니다. 이를 통해 필수적인 공간적 또는 문맥적 정보는 보존하면서 불필요한 계산 부하는 제거하여, PyTorch와 같은 프레임워크가 복잡한 비전 모델을 훨씬 더 빠르게 처리할 수 있게 합니다.

Link to this section토큰 병합의 실제 적용 사례#

토큰 병합은 계산 자원이 제한된 환경에서 무거운 어텐션 기반 아키텍처를 배포하기 위한 핵심 최적화 전략이 되었습니다.

1. Generative AI and Image Synthesis: 널리 사용되는 텍스트-이미지 확산(diffusion) 모델에서 ToMe는 이미지 생성을 가속하는 데 자주 사용됩니다. 배경이나 상세 정보가 적은 토큰들을 병합함으로써 생성 과정에 필요한 단계가 줄어들어, 엄청난 GPU 자원을 절약하고 생성형 모델을 사용하는 최종 사용자의 대기 시간을 줄여줍니다. 확산 프로세스에 대한 자세한 내용은 arXiv의 기초 연구에서 확인할 수 있습니다.

2. Edge AI Deployments: Segment Anything Model (SAM)과 같이 거대한 모델을 모바일 장치에 배포하는 것은 메모리 제한으로 인해 매우 어렵습니다. ToMe는 메모리 사용량을 동적으로 줄여 복잡한 image segmentation 작업을 에지 하드웨어에서 실행할 수 있게 돕습니다. 순수 속도가 중요한 시나리오에서 엔지니어들은 더 빠르고 완벽한 에지 추론을 위해 Ultralytics YOLO26과 같이 기본적으로 최적화된 어텐션 없는 아키텍처로 전환하기도 합니다.

Link to this sectionPython 예제: 토큰 유사도 계산#

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을 통해 크게 간소화됩니다. 이 플랫폼은 자동화된 data annotation, 원활한 클라우드 학습, 그리고 다양한 edge computing 하드웨어 환경 전반에 걸친 강력한 모델 배포를 위한 직관적인 생태계를 제공합니다. computer vision 이니셔티브를 확장하는 조직들은 이러한 도구를 활용하여 최첨단 모델을 안정적이고 효율적으로 프로덕션 환경에 배포합니다.