Token Merging (ToMe)

Слияние токенов (ToMe) — метод ускорения трансформеров, особенно визуального трансформера (ViT): похожие токены объединяются внутри сети, и каждый блок обрабатывает меньше токенов. Изображения содержат много избыточной информации: например, фрагменты неба или дороги могут почти не отличаться друг от друга. Их усреднение сокращает вычисления, сохраняя содержащуюся в них информацию. ToMe работает с уже обученными моделями без дообучения и может применяться и во время обучения. Дэниел Болья и его коллеги из Georgia Tech и Meta AI представили этот метод в работе Token Merging: Your ViT But Faster, доложенной на ICLR 2023. Без дополнительной настройки метод удвоил пропускную способность ViT-L при разрешении 512 пикселей и ViT-H при разрешении 518 пикселей на изображениях, снизив точность всего на 0,2–0,3%, а также ускорил видеомодель ViT-L в 2,2 раза при сопоставимом снижении точности.

Как работает слияние токенов#

ViT разбивает изображение на патчи, каждому из которых соответствует один токен, а каждый блок выполняет самовнимание, вычислительная стоимость которого растёт пропорционально квадрату числа токенов, а затем MLP, стоимость которой растёт линейно. ToMe добавляет шаг слияния между вниманием и MLP в каждом блоке и каждый раз удаляет фиксированное число токенов r, поэтому в 24-блочной модели с r = 8 в конце будет на 192 токена меньше, чем в начале. Для слияния используется мягкое двудольное сопоставление, выбранное за быстроту и отсутствие необходимости в итеративной кластеризации:

1. Разбиение: Поочерёдно раздели токены на два набора: A и B.
2. Сопоставление: Свяжи каждый токен из A с наиболее похожим токеном из B, измеряя сходство по косинусному сходству ключей внимания, усреднённых по головам.
3. Выбор лучших рёбер: Оставь r наиболее похожих пар.
4. Слияние: Усредни каждую выбранную пару, превратив её в один токен; взвешивай усреднение по числу исходных патчей, которые уже представляет каждый токен, а затем снова объедини два набора.

Поскольку объединённые токены представляют несколько патчей, ToMe добавляет логарифм размера каждого токена к оценкам внимания. Этот приём называется пропорциональным вниманием и позволяет токену, составленному из множества патчей, получать соответствующую долю внимания. Вся процедура дифференцируема, поэтому модели можно обучать с ToMe для дополнительного ускорения.

В этом примере на PyTorch применяется мягкое двудольное сопоставление к ViT DINOv2-small от Meta из библиотеки Transformers от Hugging Face на изображении автобуса размером 448 x 448; после каждого из 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 токенов на блок снижает его до 673 при сходстве 0.887. Это показывает, как параметр r задаёт компромисс между вычислительными затратами и точностью.

Где применяется слияние токенов#

Метод Token Merging for Fast Stable Diffusion применяет ту же идею к блокам трансформера сети удаления шума Stable Diffusion: он может удалять до 60% токенов, генерируя изображения вдвое быстрее и используя до 5,6 раза меньше памяти без дополнительного обучения. Помимо изображений, в работе о ToMe пропускная способность аудиомодели ViT-B, обученной со слиянием токенов, удвоилась при снижении mAP на 0,4%, а дообучение маскированного автоэнкодера на видео ускорилось до 2 раз. Исследователи применяют похожие методы слияния и прореживания к визуальным токенам визуально-языковых моделей, в которых LLaVA-1.5 выделяет 576 токенов контекста языковой модели для каждого изображения.

В пакете Ultralytics модель Meta Segment Anything (SAM) кодирует каждое изображение размером 1024 x 1024 пикселя в виде 4 096 токенов патчей. В сравнении с MobileSAM энкодер изображений ViT-H занимает 452 мс из 456 мс, которые SAM тратит на одно изображение при работе на GPU, — именно на эту долю вычислений и нацелены методы сокращения числа токенов.

Слияние токенов и связанные понятия#

• Прореживание токенов: Такие методы, как DynamicViT и EViT, отбрасывают токены, признанные неважными, теряя содержащуюся в них информацию и обычно требуя обучения модели. ToMe вместо этого объединяет токены, усредняя их информацию, и работает с моделями, которые не обучались с применением этого метода.
• Токенизация: Создаёт исходные токены из патчей или текста до запуска сети. ToMe сокращает их число внутри сети.
• Flash attention и разреженное внимание: Снижают стоимость или ускоряют вычисление внимания при заданном числе токенов. ToMe сокращает число токенов, ускоряя также слои MLP, и может сочетаться с этими методами.
• Смесь глубин (MoD): Позволяет обучаемому маршрутизатору решать, какие токены пропустят блок в языковых моделях. ToMe необратимо объединяет избыточные токены по фиксированному правилу сходства.
• Снижение входного разрешения: Равномерно сокращает число токенов по всему изображению. ToMe объединяет токены главным образом в избыточных областях, например на фоне, сохраняя полное разрешение детализированных областей.