Token Merging (ToMe)

Token Merging (ToMe) — это передовая методика, разработанная для оптимизации производительности и эффективности архитектур Transformer за счет сокращения количества токенов, обрабатываемых во время прямого прохода. Изначально созданный для ускорения моделей Vision Transformer (ViT), ToMe работает путем систематического поиска и объединения избыточных токенов внутри сети без необходимости дополнительного обучения. Поскольку вычислительная сложность механизма self-attention растет квадратично по отношению к количеству токенов, объединение похожих токенов значительно снижает общее количество операций с плавающей запятой (FLOPs), обеспечивая существенно более быстрый real-time inference.

Link to this sectionПонимание процесса Token Merging#

ToMe принципиально отличается от tokenization, которая является начальным этапом предварительной обработки, заключающимся в разбиении изображения или текста на отдельные tokens. В то время как токенизация создает дискретные элементы, Token Merging выступает в роли механизма динамического понижающего сэмплирования во время прямого выполнения модели.

Алгоритм обычно использует двудольное сопоставление (bipartite matching) для оценки сходства токенов, часто вычисляя косинусное сходство между ключами токенов в слоях внимания. Токены, содержащие очень похожую визуальную или семантическую информацию, объединяются — часто путем усреднения их признаков. Это гарантирует, что важная пространственная или контекстная информация сохраняется при исключении ненужной вычислительной нагрузки, что позволяет таким фреймворкам, как PyTorch, обрабатывать сложные модели компьютерного зрения гораздо быстрее.

Link to this sectionПрактическое применение Token Merging#

Token Merging стал критически важной стратегией оптимизации для развертывания тяжелых архитектур на основе внимания в вычислительно ограниченных средах.

1. Generative AI and Image Synthesis: В популярных диффузионных моделях преобразования текста в изображение ToMe часто используется для ускорения генерации изображений. За счет объединения токенов фона или токенов с низким уровнем детализации процесс генерации требует меньше шагов, что экономит огромные ресурсы GPU и снижает задержку для конечных пользователей, полагающихся на генеративные модели. Ты можешь узнать больше о диффузионных процессах в фундаментальных исследованиях на arXiv.

2. Edge AI Deployments: Развертывание массивных моделей, таких как Segment Anything Model (SAM), на мобильных устройствах крайне затруднительно из-за ограничений по памяти. ToMe помогает динамически уменьшать объем используемой памяти, позволяя запускать сложные задачи image segmentation на граничных устройствах. Для сценариев, где критически важна чистая скорость, инженеры часто переключаются на нативно оптимизированные архитектуры без внимания, такие как Ultralytics YOLO26, для более быстрого сквозного инференса на граничных устройствах.

Link to this sectionПример на 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)

Современные конвейеры машинного обучения требуют тщательного баланса точности и скорости. Независимо от того, используешь ли ты Token Merging для оптимизации кастомной модели ViT или полагаешься на передовую эффективность YOLO26, управление этими сложными рабочими процессами данных значительно упрощается с помощью Ultralytics Platform. Платформа предоставляет интуитивно понятную экосистему для автоматизированной data annotation, бесшовного облачного обучения и надежного развертывания моделей в различных средах edge computing. Организации, масштабирующие свои инициативы в области computer vision, полагаются на эти инструменты для надежного и эффективного внедрения современных моделей в производство.