Token Merging (ToMe)

دمج الرموز (ToMe) طريقة لتسريع المحوّلات، ولا سيما محوّل الرؤية (ViT)، عن طريق دمج الرموز المتشابهة داخل الشبكة، بحيث تعالج كل كتلة عددًا أقل منها. تحتوي الصور على قدر كبير من التكرار، مثل رقع السماء أو الطريق التي تبدو متشابهة جدًا، ويؤدي حساب متوسط هذه الرموز ودمجها إلى تقليل العمليات الحسابية مع الحفاظ على معلوماتها. يعمل ToMe على النماذج المدرّبة مسبقًا من دون إعادة تدريب، ويمكن تطبيقه أيضًا أثناء التدريب. وقد قدّمه دانيال بوليا وزملاؤه في Georgia Tech وMeta AI في Token Merging: Your ViT But Faster، وهي ورقة قُدّمت شفهيًا في ICLR 2023. وباستخدامه مباشرة، ضاعف ToMe معدل إنتاجية 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 التابع لـMeta، والمتاح عبر Hugging Face Transformers، وذلك على صورة حافلة مقاسها 448 x 448، مع دمج 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 صوتي دُرّب باستخدام الدمج، مع انخفاض قدره 0.4% في mAP، وسرّعت الضبط الدقيق لـالمشفّر التلقائي المقنّع على الفيديو بما يصل إلى ضعفين. ويطبق الباحثون أساليب دمج وتشذيب مشابهة على الرموز المرئية في نماذج الرؤية واللغة، حيث يستهلك LLaVA-1.5 عدد 576 رمزًا من سياق النموذج اللغوي لكل صورة.

في حزمة Ultralytics، يرمّز نموذج تجزئة أي شيء (SAM) من Meta كل صورة مقاسها 1024 x 1024 بكسل باستخدام 4,096 رمز رقعة، وفي مقارنة MobileSAM يستغرق مُرمّز الصور ViT-H مدة 452 ms من إجمالي 456 ms التي يحتاجها SAM لكل صورة على GPU، وهو الجزء الأكبر من العمليات الحسابية الذي تستهدفه أساليب تقليل الرموز.

دمج الرموز والمفاهيم ذات الصلة#

• تشذيب الرموز: تحذف أساليب مثل DynamicViT وEViT الرموز التي ترى أنها غير مهمة، ما يؤدي إلى فقدان معلوماتها ويتطلب عادةً التدريب. أما ToMe فيدمج الرموز بدلًا من حذفها، لذلك تُحفظ معلوماتها في المتوسط، ويعمل على النماذج التي دُرّبت من دونه.
• تقطيع الرموز: ينشئ الرموز الأولية من الرقع أو النص قبل تشغيل الشبكة. أما ToMe فيقلل عددها داخل الشبكة.
• انتباه Flash والانتباه المتناثر: يجعلان الانتباه أقل كلفة أو أسرع عند عدد ثابت من الرموز. أما ToMe فيقلل عدد الرموز، ما يسرّع طبقات MLP أيضًا، ويمكن دمجه معهما.
• مزيج الأعماق (MoD): يتيح لموجّه متعلّم تحديد الرموز التي تتجاوز كتلة في النماذج اللغوية. أما ToMe فيدمج الرموز الزائدة دمجًا دائمًا باستخدام قاعدة تشابه ثابتة.
• خفض دقة الإدخال: يقلل عدد الرموز بالتساوي عبر الصورة. أما ToMe فيدمج الرموز أساسًا في المناطق المتكررة، مثل الخلفية، ويحافظ على المناطق الغنية بالتفاصيل بدقتها الكاملة.