آموزش ساخت لایه Patch Embedding و ورودی ترنسفورمر بینایی (ViT)

مدل‌های ترنسفورمر در ابتدا برای پردازش متن و دنباله‌ای از توکن‌ها طراحی شدند، اما چالش اصلی در پردازش تصویر، حجم سرسام‌آور داده‌های پیکسلی است. اگر بخواهیم تک‌تک پیکسل‌های یک تصویر معمولی را مانند کلمات متن توکن در نظر بگیریم، پیچیدگی لایه‌های توجه (Attention) به‌شدت افزایش می‌یابد و بار محاسباتی مدل از توان پردازشی فراتر می‌رود.

راهکار بنیادین معماری ترنسفورمر بینایی (Vision Transformer یا ViT)، استفاده از امبدینگ تکه‌ها (Patch Embedding) به عنوان یک توکن‌ساز برای تصاویر است. در این شیوه، تصویر به شبکه‌ای منظم از تکه‌های مربعی برش داده می‌شود و هر تکه از طریق یک تصویرسازی خطی به برداری با بعد پنهان مدل تبدیل می‌گردد. در این درس با نحوه‌ی پیاده‌سازی استاندارد و سریع این لایه با استفاده از Conv2d و افزودن امبدینگ موقعیت دوبعدی آشنا می‌شوید.

یک مدل بینایی ماشین که پیکسل‌ها را می‌خواند، به یک توکن‌ساز (tokenizer) برای پیکسل‌ها نیاز دارد. امبدینگ تکه‌ها (patch embedding) دقیقاً همین توکن‌ساز است. تصویر را به شبکه‌ای از مربع‌ها برش بزنید، هر مربع را به یک بردار تخت تبدیل کنید، آن را با یک لایه‌ی خطی تصویر کنید و در نهایت یک سیگنال موقعیت دوبعدی به آن بیفزایید تا ترنسفورمر بداند هر مربع در کجای تصویر اصلی قرار داشته است.

اهداف یادگیری

• توکن‌سازی تصویر و تبدیل آن به دنباله‌ای با طول ثابت از امبدینگ تکه‌ها.
• پیاده‌سازی تصویر کردن تکه‌ها با Conv2d به‌گونه‌ای که از نظر ریاضی دقیقاً با روش «unfold و سپس لایه‌ی خطی» یکسان باشد.
• ساخت یک امبدینگ موقعیت سینوسی دوبعدی معین (deterministic) تا ترتیب توکن‌ها موقعیت مکانی را کدگذاری کند.
• اعتبارسنجی تعداد تکه‌ها، ابعاد امبدینگ و برابری خروجی Conv2d با unfold روی داده‌های آزمایشی ساختگی (fixture).

صورت مسئله

ترنسفورمر دنباله‌ای از بردارها را پردازش می‌کند، درحالی‌که تصویر یک شبکه‌ی 3 کاناله است. خواندن تک‌تک پیکسل‌ها به عنوان یک توکن، طول دنباله را به‌شدت منفجر می‌کند: یک تصویر RGB با ابعاد 224x224 شامل 150,528 توکن است؛ هزینه‌ی محاسباتی اتنشن برای چنین طولی در یک ترنسفورمر 12 لایه سرسام‌آور خواهد بود. از طرف دیگر، تبدیل کل تصویر به یک بردار تخت و غول‌پیکر، اطلاعات مجاورت مکانی (locality) را کاملاً دور می‌ریزد؛ چیزی که لایه‌ی اتنشن هرگز نمی‌تواند آن را بازیابی کند. وظیفه‌ی بخش ابتدایی انکودر، فشرده‌سازی شبکه‌ی پیکسلی به چندصد توکن است که هرکدام چکیده‌ای از یک ناحیه‌ی مربعی باشند.

امبدینگ تکه‌ها این چالش را با یک تصویر کردن خطی حل می‌کند. اگر یک تصویر 224x224 را به تکه‌های 16x16 تقسیم کنید، یک شبکه‌ی 14x14 شامل 196 تکه به دست می‌آید. هر تکه از ابعاد (3, 16, 16) = 768 مقدار پیکسلی به یک بردار تخت تبدیل می‌شود؛ سپس یک لایه‌ی خطی آن را به بعد پنهان مدل تصویر می‌کند. ترنسفورمر 196 توکن با بعد hidden (معمولاً 768) به‌همراه یک توکن CLS دریافت می‌کند. پردازش چنین دنباله‌ای برای ادامه‌ی شبکه کاملاً امکان‌پذیر و سبک است.

مفاهیم پایه

چرا تکه‌ها، نه پیکسل‌ها؟

پیچیدگی محاسباتی اتنشن بر حسب طول دنباله، درجه‌دو (quadratic) است. هزینه‌ی یک دنباله‌ی 196 توکنی، 196 * 196 = 38,416 محاسبه‌ی امتیاز اتنشن به ازای هر سر در هر لایه است؛ درحالی‌که یک دنباله‌ی 150,528 توکنی به 150,528 * 150,528 = 22.6 billion محاسبه نیاز دارد. استفاده از تکه‌ها بار محاسباتی اتنشن را 590,000 برابر کاهش می‌دهد، ضمن اینکه یک ناحیه‌ی 16x16 سیگنال کافی برای وظایف سطح بالای بینایی را در خود دارد. بهای این کار از دست رفتن جزئیات دقیق مکانی درون هر تکه است؛ به همین دلیل در پایپ‌لاین‌های چندحالته (multimodal)، هر جا مکان‌یابی دقیق اهمیت پیدا کند، معمولاً از یک branch با وضوح تصویر بالا در کنار آن استفاده می‌کنند.

چرا یک تصویر کردن خطی کافی است؟

با هر تکه مانند یک بردار مستقل رفتار می‌شود. عمل تصویر کردن، یک پایه (basis) یاد می‌گیرد: آشکارسازهای لبه، فیلترهای رنگی و بافت‌های ساده. یک لایه‌ی خطی تکی بسیار کوچک است (589,824 پارامتر برای ViT-Base با ابعاد 768 * 768) و با سرعت بالایی آموزش می‌بیند. ساقه‌های کانولوشنیِ عمیق‌تر نیز وجود دارند (مدل‌های موسوم به ViT هیبریدی)، اما تصویر کردن خطیِ تخت استاندارد اصلی است و بیشتر انکودرهای متن‌باز امروزی دقیقاً با همین ساختار عرضه می‌شوند.

ترفند Conv2d

یک لایه‌ی Conv2d(in_channels=3, out_channels=hidden, kernel_size=patch_size, stride=patch_size) بدون پدینگ (padding)، از نظر عددی نتیجه‌ای کاملاً یکسان با روش unfold و لایه‌ی خطی می‌دهد؛ چون هر موقعیت خروجی، در واقع ضرب داخلی پیکسل‌های تکه در یک فیلتر است. همین کانولوشن کار تصویر کردن تکه را انجام می‌دهد و بیشتر کدهای محیط عملیاتی (پروداکشن) هم از همین روش استفاده می‌کنند، زیرا روی GPU سریع‌تر است و به یک تغییر شکل (reshape) کمتر نیاز دارد.

امبدینگ‌های موقعیت

توکن‌ها پس از خروج از لایه‌ی تصویر، هیچ ترتیب مکانی مشخصی ندارند. امبدینگ سینوسی دوبعدی سیگنال ثابتی به هر توکن اختصاص می‌دهد که موقعیت (row, col) آن را کدگذاری می‌کند. نیمی از بعد امبدینگ، موقعیت سطر را با استفاده از sin/cos در فرکانس‌های مختلف کدگذاری می‌کند و نیم دیگر، موقعیت ستون را پوشش می‌دهد. این روش کدگذاری کاملاً معین و قطعی است؛ بنابراین می‌توانید بدون نیاز به بازآموزی مدل، وضوح تصاویر ورودی را تغییر دهید و مقادیر را حتی برای ابعادی که مدل حین آموزش ندیده، به‌زیبایی درون‌یابی (interpolate) کنید.

برای مدل ViT-Base/16 با وضوح 224: لایه‌ی تصویر شامل 590,592 پارامتر است، توکن CLS دارای 768 پارامتر بوده و امبدینگ سینوسی موقعیت هیچ پارامتری مصرف نمی‌کند. در درس بعدی (درس 59)، یک ترنسفورمر 12 لایه را روی همین بخش ابتدایی سوار می‌کنیم.

بررسی هم‌ارزی برای اطمینان از درستی کار

مرحله‌ی خرد کردن تکه‌ها را به دو شکل می‌توان نوشت: تصویر کردن با Conv2d یا ترکیب صریحِ unfold و سپس لایه‌ی خطی. هر دو روش در ازای وزن‌های یکسان باید خروجی کاملاً برابری تولید کنند. اگر خروجی‌ها یکی نباشد، منطق ریاضی unfold اشتباه است و بقیه‌ی انکودر روی پایه‌ای سست بنا خواهد شد. تست‌های این درس درستی این هم‌ارزی را می‌سنجند.

پیاده‌سازی از صفر

فایل code/main.py این بخش‌ها را پیاده‌سازی می‌کند:

• کلاس PatchEmbed، یک nn.Module که Conv2d را برای تصویر کردن تکه‌ها کپسوله می‌کند.
• تابع sinusoidal_2d(grid_h, grid_w, dim)، تابعی بی‌وضعیت (stateless) که جدول موقعیت‌های دوبعدی را می‌سازد.
• کلاس VisionFrontEnd، که امبدینگ تکه‌ها، افزودن توکن CLS به ابتدای دنباله و جمع امبدینگ موقعیت را در یک گذر رو به جلو (forward pass) ترکیب می‌کند.
• تابع کمکی synthesize_image(seed) که داده‌های آزمایشی معین در ابعاد 224x224x3 را با numpy.random می‌سازد.
• یک نمونه‌ی اجرایی (demo) که یک تصویر آزمایشی را از این بخش عبور می‌دهد و ابعاد خروجی، نُرم توکن CLS و یک سطر از امبدینگ موقعیت را چاپ می‌کند.

کد را اجرا کنید:

python3 code/main.py

خروجی: تصویر آزمایشی 224x224 به دنباله‌ای با ابعاد (1, 197, 768) توکن‌سازی می‌شود. اولین توکن، CLS است و 196 توکن بعدی، توکن‌های وصله‌ای (patch tokens) هستند. نرم امبدینگ‌های موقعیت درون یک سطر کاملاً یکنواخت است که این امضای توابع سینوسی به شمار می‌رود.

استفاده از ابزارهای واقعی

همین بخش ابتداییِ تکه‌بندی تصویر در تمام مدل‌های مدرن بینایی-زبان دیده می‌شود: CLIP ViT-L/14، SigLIP، DINOv2، خانواده‌ی Qwen-VL و InternVL، همگی کارشان را با تصویر کردن تکه‌ها از طریق Conv2d و افزودن سیگنال موقعیت شروع می‌کنند. تفاوت این خانواده‌ها در لایه‌های بعدی نهفته است (استخرسازی با CLS در مقایسه با بدون CLS، توکن‌های رجیستر (register tokens)، اندازه‌ی متفاوت تکه‌ها مثل 14 در برابر 16، و وضوح تصویر پویا از طریق درون‌یابی موقعیت‌ها). ساختاری که در این درس می‌سازید، پایه‌ای است که تمام این مدل‌ها روی آن بنا شده‌اند.

تست‌ها

فایل code/test_main.py این موارد را پوشش می‌دهد:

• مطابقت تعداد تکه‌ها با رابطه‌ی (image_size / patch_size) ** 2
• تطابق ابعاد خروجی با (batch, num_patches + 1, hidden)
• برابری لایه‌ی تصویر Conv2d با روش دستی unfold و لایه‌ی خطی روی داده‌های آزمایشی کوچک
• یکسان و معین بودن خروجی جدول موقعیت سینوسی در فراخوانی‌های مکرر
• انتشار (broadcast) درست توکن CLS در بعد batch بدون هیچ نشت راه‌حل (leakage)

تست‌ها را اجرا کنید:

python3 -m unittest code/test_main.py

تمرین‌ها

1. امبدینگ موقعیت سینوسی را با یک nn.Parameter یادگرفتنی جایگزین کنید و تابع زیان (loss) دور اول آموزش (epoch 1) را در یک وظیفه (task) دسته‌بندی ساختگی کوچک بسنجید. در وضوح ثابت، موقعیت‌های یادگرفتنی بهتر عمل می‌کنند؛ اما اگر بعد از آموزش وضوح تصویر را تغییر دهید، روش سینوسی برنده است.

2. لایه‌ی Conv2d را با ترکیب صریح nn.Unfold و nn.Linear جایگزین کنید و با دستور assert بررسی کنید که آیا خروجی‌ها تا حد خطای محاسبات اعشاری (float tolerance) یکسان هستند یا خیر. منطق ریاضی هر دو یکی است و فقط نحوه‌ی پیاده‌سازی فرق دارد.

3. پشتیبانی از تکه‌های غیرمربعی (مثلاً 32x16 برای ورودی‌های عریض) را اضافه کنید و مطمئن شوید جدول موقعیت از پس شبکه‌های غیرمربعی هم برمی‌آید.

4. زمان اجرای مرحله‌ی خرد کردن تکه‌ها را در اندازه‌های batch معادل 1، 8 و 64 پروفایل کنید. تصویر کردن تکه‌ها به‌ندرت گلوگاه سیستم می‌شود؛ بخش عمده‌ی زمان صرف لایه‌های اتنشن بعدی خواهد شد.

5. این بخش ابتدایی را به عنوان یک استخراج‌کننده‌ی ویژگیِ منجمدشده (frozen) روی یک مجموعه‌داده‌ی ساختگی با 4 کلاس از شکل‌های هندسی (دایره، مربع، مثلث، ستاره) آموزش دهید. خروجی توکن CLS باید این کلاس‌ها را به‌صورت خطی تفکیک کند.

اصطلاحات کلیدی

برای مطالعه‌ی بیشتر

• مقاله‌ی An Image is Worth 16x16 Words (معرفی ViT، سال 2021) برای آشنایی با ایده‌ی اصلی امبدینگ تکه‌ها.
• مقاله‌ی Attention Is All You Need (سال 2017) برای فرمول موقعیت سینوسی که در اینجا به حالت دوبعدی تعمیم یافته است.
• مقاله‌ی DINOv2 برای مطالعه‌ی توکن‌های رجیستر (register tokens)؛ قابلیتی که می‌توانید آن را به عنوان تمرین ششم پیاده‌سازی کنید.