راهنمای فنی Kauldron؛ ابزار Google Research برای آموزش مدل‌های JAX

کتابخانه Kauldron در راهنمای تازهٔ MarkTechPost به‌عنوان ابزاری از Google Research برای آموزش مدل‌های مبتنی بر JAX بررسی شده است؛ راهنما به‌جای معرفی یک مدل یا اعلام عرضهٔ تازه، نشان می‌دهد این کتابخانه چگونه پیکربندی آزمایش‌ها، اتصال اجزای نرم‌افزاری و اجرای آموزش را سامان می‌دهد. اهمیت این بررسی در نمایش جزئیات عملی یک ابزار پژوهشی است: نویسنده یک مدل را با دادهٔ مصنوعی آموزش می‌دهد، چند نسخهٔ آزمایشی می‌سازد و نشان می‌دهد آموزش پس از توقف از نقطهٔ ذخیره‌شده ادامه پیدا می‌کند.

این راهنما Kauldron را در برابر مجموعه‌ای از ابزارهای معمول در جریان کار JAX قرار می‌دهد و سه سازوکار اصلی را توضیح می‌دهد: konfig برای تعریف و تغییر پیکربندی، kontext برای متصل‌کردن اجزا از راه مسیرهای متنی، و بررسی شکل آرایه‌ها در زمان اجرا. سپس همین اجزا در یک نمونهٔ آموزشی به کار گرفته می‌شوند. طبق گزارش نویسنده، نمونه روی پردازندهٔ مرکزی اجرا شده و برای داده‌ها نیازی به دانلود مجموعه‌داده یا استفاده از شتاب‌دهنده نداشته است.

کتابخانه Kauldron و پیکربندی آزمایش‌ها به‌صورت داده

JAX کتابخانه‌ای برای محاسبات عددی و یادگیری ماشین است که در این نمونه، همراه با Optax برای بهینه‌سازی و اجزای دیگر جریان آموزش استفاده می‌شود. Kauldron لایه‌ای برای تعریف و اجرای این فرایند فراهم می‌کند. محور نخست این طراحی، konfig است: به‌جای آن‌که پیکربندی یک آزمایش از ابتدا به مجموعه‌ای از اشیای آماده‌شده تبدیل شود، اطلاعات آن به شکل درختی از مقادیر و فراخوانی‌های توصیف‌شده نگهداری می‌شود.

راهنما مثالی از ساخت بهینه‌ساز Optax با نرخ یادگیری مشخص ارائه می‌کند. در این مرحله، پیکربندی هنوز خودِ بهینه‌ساز نیست؛ نام تابع و آرگومان‌های آن را نگه می‌دارد و تنها هنگام resolve به شیء واقعی تبدیل می‌شود. از آن‌جا که تنظیمات در قالب دادهٔ تو‌در‌تو باقی می‌مانند، می‌توان آن‌ها را به JSON تبدیل کرد و بعد دوباره به پیکربندی قابل‌اجرا برگرداند. به گفتهٔ نویسنده، این روش به Optax نیاز ندارد از پیش برای Kauldron بازنویسی یا مجهز شده باشد.

یک ویژگی دیگر، ارجاع‌دادن یک مقدار به تنظیم دیگری در پیکربندی است. در نمونهٔ مقاله، تعداد گام‌های آموزش در تعیین گام‌های کاهش نرخ یادگیری به کار می‌رود. اگر پژوهشگر تعداد گام‌ها را تغییر دهد، زمان‌بندی نرخ یادگیری نیز با آن هماهنگ می‌شود. این پیوند مانع از آن می‌شود که یک مقدار قدیمی در بخشی از تنظیمات باقی بماند و نتیجهٔ مقایسهٔ آزمایش‌ها را مخدوش کند.

این شکل از پیکربندی به اجرای چند آزمایش با تغییرات محدود هم کمک می‌کند. نویسنده پنج نسخه را با تغییر یک خط در هرکدام اجرا کرده است؛ تغییرها شامل پهنای مدل و تنظیم بهینه‌ساز بوده و در یکی از نسخه‌ها Adam با SGD جایگزین شده است. بر اساس راهنما، بدنهٔ مدل، تابع خطا و حلقهٔ آموزش میان این نسخه‌ها ثابت مانده‌اند.

اتصال اجزا با مسیرهای متنی و بررسی شکل آرایه‌ها

سازوکار دوم، kontext، برای رساندن دادهٔ موردنیاز به اجزای مختلف استفاده می‌شود. یک مسیر متنی مانند batch.image می‌تواند به بخشی از ساختار داده اشاره کند؛ مسیرهای این‌چنینی به کلیدهای دیکشنری، ویژگی‌های شیء یا جایگاه‌های فهرست دسترسی پیدا می‌کنند. اگر مسیر وجود نداشته باشد، سامانه خطایی می‌دهد که به گفتهٔ نویسنده، اطلاعات مربوط به داده‌های در دسترس را نیز نشان می‌دهد.

در نمونهٔ راهنما، معیار ارزیابی ورودی‌هایش را با چنین مسیرهایی اعلام می‌کند و Kauldron مقادیر مربوط را به آن می‌دهد. در نتیجه، معیار لازم نیست مدل را وارد کند یا از ساختار داخلی آن اطلاع داشته باشد. اگر لازم باشد معیار به خروجی دیگری نگاه کند، تغییر مسیر متنی کافی است. این روش اجزای آزمایش را از وابستگی مستقیم به یکدیگر دور نگه می‌دارد و امکان جایگزین‌کردن بخشی از جریان کار را ساده‌تر می‌کند.

قابلیت دیگر به بررسی شکل آرایه‌ها مربوط است. شکل یا ابعاد یک آرایه مشخص می‌کند داده چه اندازه و ساختاری دارد؛ ناسازگاری ابعاد میان ورودی‌ها می‌تواند باعث خطا در محاسبات شود. Kauldron در نمونهٔ مقاله اجازه می‌دهد محورهایی مانند c در امضای تابع نام‌گذاری شوند و هنگام اجرا مقدارشان بررسی شود. اگر دو آرایه در محوری که باید مشترک باشد اندازه‌های متفاوت داشته باشند، خطا علاوه بر شکل آرایه‌ها توضیح می‌دهد آن محور پیش‌تر به چه مقداری مقید شده است.

به گزارش نویسنده، این جزئیات می‌تواند تشخیص خطا را در مدلی که چندین آرایه در بخش‌های مختلف آن جریان دارند، آسان‌تر کند. به‌جای مقایسهٔ چند مجموعه عدد بدون نام، پژوهشگر می‌بیند اختلاف به کدام محور مربوط است. مقاله این بررسی را در زمان اجرا انجام می‌دهد؛ یعنی ناسازگاری هنگام اجرای برنامه آشکار می‌شود، نه صرفاً با خواندن تعریف تابع.

آموزش مدل، محاسبهٔ معیار و اجرای آزمایش‌ها

برای نشان‌دادن مسیر کامل آموزش، نویسنده یک تابع خطا و یک معیار سفارشی می‌سازد. هر دو در قالبی تعریف می‌شوند که Kauldron انتظار دارد و ورودی‌هایشان را با کلیدهای kontext دریافت می‌کنند. تابع خطا مقدارهایی را برای هر جزء داده برمی‌گرداند و چارچوب کاهش آن‌ها را مدیریت می‌کند. راهنما می‌گوید در نمونهٔ اجرا، وزن نیم باعث شده مقدار نهایی تابع خطا نصف شود.

معیار سفارشی به‌جای بازگرداندن یک عدد مستقل، حالتی را نگه می‌دارد که می‌توان آن را با حالت‌های دیگر ادغام کرد. نمونه دو مقدار جمع‌شونده، صورت و مخرج، را برای محاسبه به کار می‌گیرد. نویسنده یک دستهٔ شش‌ردیفی و دستهٔ دیگری با دو ردیف را ترکیب می‌کند و نشان می‌دهد میانگین وزنی کل با میانگین سادهٔ دو نرخ دسته‌ای یکسان نیست. این تمایز برای دسته‌های نامساوی اهمیت دارد؛ در راهنما همچنین آمده است که سازوکار ادغام می‌تواند نتیجه‌ها را از چند دستگاه جمع کند، مستقل از ترتیب رسیدن آن‌ها.

برای آموزش، دادهٔ رگرسیون مصنوعی در حافظه ساخته می‌شود و به خط لوله‌ای داده می‌شود که دسته‌بندی و درهم‌ریزی داده‌ها را انجام می‌دهد. سپس مدل، داده، تابع خطا، معیار و بهینه‌ساز در Trainer کنار هم قرار می‌گیرند. طبق نتیجهٔ گزارش‌شده در مقاله، مقدار خطا در ۳۰۰ گام از ۱٫۵۱ به ۰٫۰۰۵ کاهش یافته و اجرا روی پردازندهٔ مرکزی حدود یک ثانیه طول کشیده است. این نتیجه به نمونهٔ کوچک و دادهٔ مصنوعی مقاله مربوط است و به‌تنهایی دربارهٔ عملکرد مدل‌ها یا سرعت آموزش در کاربردهای دیگر چیزی ثابت نمی‌کند.

Trainer همچنین اجازه می‌دهد خروجی‌های جانبی، مانند مقادیر تابع خطا و معیارها، فقط در صورت درخواست محاسبه شوند؛ راهنما دلیل این رفتار را هزینهٔ محاسباتی می‌داند. برای پایش یک لایهٔ میانی مدل نیز نیازی به ویرایش خود مدل نیست: مقاله یک مسیر متنی را در پیکربندی به خروجی لایهٔ Dense به نام enc وصل می‌کند تا فعال‌سازی‌های آن ثبت شوند.

ذخیرهٔ وضعیت و معنای عملی ماژولار بودن

بخش پایانی راهنما به ارزیابی و بازیابی آموزش می‌پردازد. ارزیاب در نمونه از مدل، توابع خطا و معیارهای تعریف‌شده در پیکربندی اصلی استفاده می‌کند و تنظیمات جداگانه‌ای برای داده و زمان‌بندی اجرا دارد. ابزار ذخیره‌سازی نیز وضعیت آموزش را در فاصله‌های گام مشخص ثبت می‌کند. نویسنده Trainer را دوباره با همان پوشهٔ کاری اجرا می‌کند و گزارش می‌دهد که فرایند از گام ۲۰۰ ادامه یافته، نه از صفر.

این رفتار برای کارهایی که ممکن است متوقف شوند کاربرد دارد: پس از بازسازی Trainer، وضعیت ذخیره‌شده امکان ادامهٔ کار را فراهم می‌کند. مقاله این مسیر را در مقیاسی کوچک و داخل نمونهٔ آموزشی نشان می‌دهد، نه در قالب گزارشی از اجرای یک پروژهٔ بزرگ یا استفاده در محیط عملیاتی.

راهنما همچنین یک ناسازگاری نسخه‌ای را در ترکیب مشخصی از ابزارها مطرح می‌کند. بنا بر توضیح نویسنده، JAX ۰٫۱۰٫۱ بخشی داخلی را جابه‌جا کرده و etils تا نسخهٔ ۱٫۱۴٫۰ همچنان به آن ارجاع می‌دهد؛ این مسیر هنگام بررسی نوع دادهٔ آرایه به خطا می‌رسد. مقاله برای نمونهٔ خود وصله‌ای دوخطی را شرح می‌دهد که از رابط عمومی JAX استفاده می‌کند و می‌گوید در نسخه‌های قدیمی‌تر بی‌اثر است. این توضیح به ترکیب نسخه‌های یادشده مربوط است و نباید آن را به همهٔ نصب‌های Kauldron تعمیم داد.

برداشت اصلی نویسنده از تجربهٔ کار با Kauldron این است که ماژولار بودن آن از جدا نگه‌داشتن پیکربندی، داده و اجزای مدل حاصل می‌شود؛ نه از تحمیل یک بازنویسی بزرگ بر کد موجود. با این حال، راهنما یک ارزیابی مستقل از همهٔ قابلیت‌های کتابخانه نیست: داده‌ها مصنوعی‌اند و نویسنده مشخصاً بخش‌هایی مانند مجموعه‌داده‌های واقعی، توزیع آموزش و ابزارهای مدیریت آزمایش را وارد نمونه نکرده است.

این خبر برای چه کسانی اهمیت دارد

این بررسی بیش از همه برای پژوهشگران یادگیری ماشین و توسعه‌دهندگانی مفید است که با JAX کار می‌کنند و می‌خواهند آزمایش‌های قابل‌تغییر و قابل‌بازیابی بسازند. برنامه‌نویسانی که کد خود را میان مدل، معیار و داده جدا می‌کنند نیز می‌توانند از ایده‌های پیکربندی به‌صورت داده، ارجاع به تنظیمات و اتصال اجزا با مسیرهای متنی استفاده کنند؛ حتی اگر تصمیمی برای به‌کارگیری Kauldron نداشته باشند.

برای خوانندگان عمومی‌تر، نکتهٔ اصلی این است که پیشرفت در پژوهش هوش مصنوعی فقط به مدل‌های تازه وابسته نیست. ابزارهایی که اجرای آزمایش‌ها را قابل‌فهم‌تر و تکرارپذیرتر می‌کنند نیز بخشی از زیرساخت این حوزه‌اند. مقالهٔ MarkTechPost نمونه‌ای عملی از این نوع ابزار ارائه می‌کند، اما خبر از عرضهٔ مدل تازه یا تغییر محصول نمی‌دهد.