🤖 machine learning

A Faster Generalized Two-Stage Approximate Top-K

تعمم هذه الورقة البحثية خوارزمية تقريبية من مرحلتين لاستخراج أفضل $K$ عنصر (Top-K) عبر اختيار أفضل $K'$ عنصر لكل قسم بدلاً من اختيار أفضل عنصر واحد فقط، مما يوفر حداً نظرياً أدق للاستدعاء (recall) ويحقق تسريعاً بمقدار رتبة عشرية على Cloud TPUv5e مع الحفاظ على نفس الاستدعاء المتوقع.

المؤلفون الأصليون: Yashas Samaga, Varun Yerram, Spandana Raj Babbula, Prateek Jain, Praneeth Netrapalli

نُشر 2026-05-14

📖 4 دقيقة قراءة☕ قراءة في استراحة قهوة

CC BY 4.0

المؤلفون الأصليون: Yashas Samaga, Varun Yerram, Spandana Raj Babbula, Prateek Jain, Praneeth Netrapalli

البحث الأصلي مرخَّص بموجب CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ هذا شرح مولَّده بالذكاء الاصطناعي للبحث أدناه. لم يكتبه المؤلفون ولم يصادقوا عليه. وللتحقق من الدقة التقنية، يرجى الرجوع إلى البحث الأصلي. اقرأ إخلاء المسؤولية الكامل

تخيل أنك مدير مكتبة ضخمة تضم ملايين الكتب (البيانات). كل يوم، يتعين عليك العثور على الكتب الـ "Top-K" الأكثر شعبية (أكبر K من الأرقام) لترشيحها للزوار.

في عالم الرقائق الحاسوبية (تحديداً تلك المستخدمة لتدريب النماذج الضخمة للذكاء الاصطنا1)، يعد العثور على هذه العناصر "الأكثر شعبية" بطيئاً ومكلفاً بشكل مفاجئ. الأمر يشبه محاولة العثور على أفضل 100 كتاب عبر قراءة كل كتاب منها واحداً تلو الآخر، رغم أن مكتبتك مصممة لإجراء العمليات الحسابية على أكوام ضخمة من الكتب دفعة واحدة.

إليك تفصيل بسيط لما يفعله هذا البحث لحل هذه المشكلة.

الطريقة القديمة: مرشح "واحد في كل مرة"

حاولت طريقة سابقة (بواسطة Chern وآخرون، 2022) تسريع ذلك باستخدام عملية من خطوتين:

التقسيم: تخيل تقسيم مكتبتك إلى 100 غرفة مختلفة (صناديق/Buckets).
المسح الأول: في كل غرفة، يقوم مساعد باختيار الكتاب الأكثر شعبية فقط ويحضره إلى مكتب الاستقبال.
الفرز النهائي: ثم ينظر المدير إلى تلك الكتب الـ 100 فقط (كتاب واحد من كل غرفة) ويختار أفضل 100 كتاب بشكل عام.

المشكلة: كانت هذه الطريقة حذرة للغاية. فمن خلال اختيار الكتاب الأفضل الوحيد من كل غرفة، غالباً ما كانت تفوتها الكتب الثانية أو الثالثة الأفضل المختبئة في نفس الغرفة. ولضمان عدم تفويت أي شيء، اضطروا لاستخدام غرف (صناديق) كثيرة جداً، مما يعني أن المدير كان لا يزال مضطراً لفرز كومة ضخمة من الكتب في النهاية. كان الأمر لا يزال بطيئاً للغاية.

الفكرة الجديدة: مرشح "Top-K"

أدرك مؤلفو هذا البحث أن الرقائق الحاسوبية تمتلك قوة إضافية لم تكن تُستخدم. لقد اقترحوا نسخة أكثر ذكاءً من الخطوة الأولى:

بدلاً من اختيار الكتاب رقم 1 فقط من كل غرفة، يقوم المساعد الآن باختيار الكتب الـ "Top-K'" (على سبيل المثال، أفضل 4 كتب) من كل غرفة.

لماذا هذا أفضل؟

عدد غرف أقل مطلوب: لأن المساعد يمسك بمزيد من الكتب من كل غرفة، فأنت لا تحتاج إلى عدد كبير من الغرف لضمان التقاط جميع الكتب الشهيرة.
فرز أقل: على الرغم من أن المساعد يمسك بمزيد من الكتب لكل غرفة، إلا أن إجمالي عدد الكتب المرسلة للمدير للفرز النهائي هو في الواقع أصغر بكثير.
النتيجة: لدى المدير كومة صغيرة لفرزها بدلاً من جبل.

"سحر" الأجهزة (Hardware)

يوضح البحث أن الرقائق الحاسوبية الحديثة (مثل Google TPU) تشبه مصانع ضخمة بها محطات عمل مختلفة:

وحدة المصفوفة (MXU): مصنع فائق السرعة يقوم بالعمليات الحسابية الثقيلة (الضرب) ولكنه سيء في الفرز.
الوحدة المتجهة (VPU): محطة عمل أصغر وأبطأ، وهي جيدة في الفرز واختيار الفائزين.

الطريقة القديمة أهدرت وقت الـ VPU. الطريقة الجديدة تستخدم الـ VPU لالتقاط كتب الـ "Top-K'" بينما تكون الـ MXU مشغولة بالقيام بالعمليات الحسابية. إنه يشبه وجود عامل يمسك بأفضل العناصر من حزام ناقل بينما لا تزال الآلة تعمل، بحيث لا يكون هناك وقت انتظار.

النتائج: تسريع الذكاء الاصطناعي

اختبر المؤلفون هذا على شريحة Google TPU:

الطريقة القديمة: استغرق العثور على الكتب الأفضل وقتاً طويلاً، وكان في كثير من الأحيان أبطأ من العمليات الحسابية التي أنشأت القائمة في المقام الأول.
الطريقة الجديدة: من خلال الحصول على "أفضل 4" من كل صندوق بدلاً من مجرد "أفضل 1"، قللوا العمل المطلوب للفرز النهائي بمعدل 7 مرات في المتوسط.
الدمج (Fusion): حتى أنهم تمكنوا من دمج خطوة "الاختيار" مع خطوة "الحساب" بحيث يحدثان في نفس الوقت تماماً.

في اختبار واقعي (العثور على أفضل 2% من البيانات في نموذج ذكاء اصطناعي كبير)، جعلت طريقتهم الجديدة العملية أسرع بـ 24 مرة من المعيار السابق. وهذا يعني أن نموذج الذكاء الاصطناعي يمكنه التدريب والعمل بشكل أسرع بكثير دون فقدان الدقة.

الخلاصة:

الطريقة القديمة: لديك 1,000 فريق. كل فريق يرسل لك أفضل لاعب لديه. ثم يتعين عليك مقابلة 1,000 لاعب لتجد أفضل 100 منهم.
الطريقة الجديدة: لديك عدد أقل من الفرق (قل 250 فريقاً). كل فريق يرسل لك أفضل 4 لاعبين لديه. سيتعين عليك مقابلة 1,000 لاعب فقط (250 فريقاً × 4 لاعبين)، ولكن لأنك حصلت على خيارات أكثر من كل فريق، فأنت على الأرجياء ستجد اللاعبين الأفضل حقاً، وتفعل ذلك بشكل أسرع بكثير لأنك نظمت الفرق بشكل أفضل.

يثبت البحث رياضياً أن نهج "Top-K'" هذا ليس مجرد تخمين؛ بل هو طريقة مضمونة للحصول على نفس جودة النتائج بعمل أقل بكثير.

ملخص تقني: خوارزمية (Top-K) التقريبية ذات المرحلتين المعممة والأسرع

بيان المشكلة
يعد تحديد العناصر الـ $K$ العليا (top- $K$ ) في مصفوفة عملية أساسية في العديد من خوارزميات تعلم الآلة، بما في ذلك البحث عن الضرب الداخلي الأقصى (MIPS)، والبحث عن أقرب $K$ جيران (KNN)، وتحسين النماذج التأسيسية الضخمة (على سبيل المثال: للحصول على التشتت في لوغاريتمات التصنيف، وكتل MLP، وآليات الانتباه، وضغط ذاكرة الكاش لـ KVCache). على المسرعات الحديثة مثل وحدات TPU وGPU، غالبًا ما يمثل حساب الـ Top- $K$ الدقيق عائقًا في الأداء. هذه الأجهزة مُحسّنة لعمليات ضرب المصفوفات الكثيفة (MXUs/TensorCores)، بينما تعتمد عمليات Top- $K$ عادةً على وحدات قياسية/متجهية (VPUs/CUDA cores) ويمكن أن تكون أبطأ بكثير من عمليات ضرب المصفوفات التي تولد البيانات. فعلى سبيل المثال، في وحدة TPUv5e، يستغرق إيجاد أعلى 2% من تنشيطات الطبقة الخفية في نموذج Gemma 2 9B زمنًا أطول بـ 27 مرة من عملية ضرب المصفوفة المقابلة لها.

بينما توفر خوارزميات الـ Top- $K$ التقريبية حلاً بديلًا، إلا أن الطريقة الأحدث المقدمة من Chern et al. (2022)، والمطبقة كـ jax.lax.approx_max_k، لا تزال تتسبب في زمن انتقال (latency) كبير. تعتمد منهجيتهم ذات المرحلتين على تقسيم المدخلات إلى حاويات (buckets)، واختيار العنصر الأول الأعلى من كل حاوية، ثم ترتيب النتائج. ومع ذلك، تظل المرحلة الثانية (ترتيب العناصر المختارة) تمثل عنق زجاجة، كما أن المرحلة الأولى غالبًا ما لا تستغل موارد الحوسبة القياسية/المتجهية المتاحة بشكل كافٍ.

المنهجية
يقترح المؤلفون تعميمًا لخوارزمية Chern et al. (2022)، حيث يقوم بتعديل المرحلة الأولى لاختيار أعلى $K'$ عنصر من كل حاوية (حيث $1 \le K' \le K$ ) بدلاً من الاقتصار على اختيار أعلى عنصر واحد فقط.

التعميم الخوارزمي:
- المرحلة الأولى: يتم تقسيم مصفوفة المدخلات بحجم $N$ إلى $B$ من الحاويات. وبدلاً من اختيار الحد الأقصى فقط، تختار الخوارزمية أعلى $K'$ عنصر من كل حاوية.
- المرحلة الثانية: يتم دمج العناصر $B \times K'$ المختارة وترتيبها لإرجاع عناصر الـ top- $K$ النهائية.
- التنفيذ: يتم تنفيذ المرحلة الأولى كبرنامج فرعي للتحديث عبر الإنترنت (online update subroutine) باستخدام Pallas (لغة نواة JAX لوحدات TPU). وهي تحافظ على قائمة مرتبة من قيم وفهارس أعلى $K'$ لكل حاوية. يتم إدراج العناصر الجديدة باستخدام تمريرة تشبه "الترتيب الفقاعي" (bubble-sort-like pass)، والمُحسّنة للعمل المتجهي عبر الحاويات. يقدم المؤلفون نسختين: "غير مدمجة" (unfused) (نواتان منفصلتان) ونسخة "مدمجة مع ضرب المصفوفات" (matmul-fused) حيث تُدمج المرحلة الأولى مع عملية ضرب المصفوفة السابقة، مما يؤدي فعليًا إلى إخفاء تكلفة المرحلة الأولى.
التحليل النظري:
- يشتق المؤلفون تعبيرًا دقيقًا للاستدعاء (recall) المتوقع لهذه الخوارزمية المعممة تحت التقسيم العشوائي، من خلال نمذجة توزيع عناصر top- $K$ في الحاويات باستخدام التوزيع الهندسي الفائق (hypergeometric distribution).
- يثبتون أنه بالنسبة لهدف استدعاء ثابت، فإن زيادة $K'$ تسمح بتقليل عدد الحاويات $B$ بحيث يتم تقليل إجمالي العناصر الممررة إلى المرحلة الثانية ( $B \times K'$ ) إلى الحد الأدنى.
- والأهم من ذلك، يشتقون حدًا جديدًا لعدد الحاويات المطلوبة عندما تكون $K'=1$ . هذا الحد هو، من الناحية المثبتة، أكثر إحكامًا بمعامل قدره 2 مقارنة بالحد في Chern et al. (2022)، مما يسمح بعدد أقل من الحاويات وبالتالي عدد أقل من العناصر لترتيبها حتى في الإعداد الأصلي.
نمذجة الأداء:
- باستخدام تحليل نقطة الحافة (ridge-point analysis)، يحدد المؤلفون الأنظمة الفرعية للمسرع (HBM، VPU، MXU). ويوضحون أن المرحلة الأولى تكون عادةً مقيدة بالذاكرة (memory-bound) حتى تصل $K'$ إلى عتبة معينة (على سبيل المثال، $K' \approx 6$ في TPUv5e). وضمن هذا النطاق المقيد بالذاكرة، لا تؤدي زيادة $K'$ إلى زيادة كبيرة في زمن انتقال المرحلة الأولى، ولكنها تقلل بشكل جذري من عبء العمل في المرحلة الثانية.

المساهمات الرئيسية

خوارزمية معممة: خوارزمية Top- $K$ تقريبية جديدة ذات مرحلتين، تختار أعلى $K'$ عنصر لكل حاوية، مما يوفر مقايضة مرنة بين حوسبة المرحلة الأولى وعبء ترتيب المرحلة الثانية.
الحدود النظرية:
- تعبير دقيق للاستدعاء المتوقع كدالة لـ $N, K, B,$ و $K'$ .
- حد أكثر إحكامًا (بمعامل 2) لعدد الحاويات المطلوبة في حالة $K'=1$ مقارنة بالعمل السابق، مما يحسن عملية اختيار المعلمات.
التنفيذ: تنفيذ عالي الأداء على Cloud TPUv5e باستخدام Pallas، يتميز بنسخ غير مدمجة ومدمجة مع ضرب المصفوفات (matmul-fused).
الاختيار التلقائي للمعلمات: إجراء يختار تلقائيًا $K'$ و $B$ لتلبية هدف استدعاء محدد من قبل المستخدم، مما يلغي الحاجة إلى الضبط اليدوي.

النتائج

الكفاءة النظرية: بالنسبة لهدف استدعاء ثابت (مثل 95%)، فإن زيادة $K'$ تقلل بشكل كبير من عدد العناصر المعالجة في المرحلة الثانية. على سبيل المثال، يتطلب اختيار أعلى 1024 من 262,144 عنصر 16,384 عنصر مخرج عند $K'=1$ ، ولكن 2,048 عنصر فقط عند $K'=4$ (انخفاض بمقدار 8 أضعاف).
تحسينات زمن الانتقال (غير المدمجة): في TPUv5e، تحقق الخوارزمية المعممة مع $K'=4$ تسارعًا بمقدار 11 ضعفًا مقارنة بـ jax.lax.approx_max_k الأصلية (27 ميكروثانية مقابل 305 ميكروثانية) لهدف استدعاء 99% لنفس المهمة.
تحسينات زمن الانتقال (المدمجة): في مهمة البحث عن الضرب الداخلي الأقصى (MIPS) (1 مليون متجه، 1024 استعلام)، تقلل نسخة التنفيذ المدمجة مع ضرب المصفوفات مع $K'=4$ إجمالي وقت التشغيل من 594 مللي ثانية (Exact Top-K) إلى 10 مللي ثانية. ويمثل هذا تسارعًا إجماليًا بمقدار 6.7 ضعفًا مقارنة بالخط المرجعي الكثيف، مما يجعل خطوة الـ Top- $K$ التقريبية أقل من تكلفة عملية ضرب المصفوفة نفسها.
التأثير النهائي: في اختبار بنية Transformer متفرقة (Gemma 2 9B)، تقلل الطريقة المقترحة من عبء كتلة MLP المتفرقة من كونها أبطأ بـ 2.7 مرة من الكتلة الكثيفة (باستخدام الخوارزمية الأصلية) إلى مجرد زيادة طفيفة قدرها +5 مللي ثانية بالنسبة للخط المرجعي الكثيف.

الأهمية والادعاءات
يزعم البحث أن اختيار $K'=1$ في خوارزميات الـ Top- $K$ التقريبية السابقة لا يستغل موارد الحوسبة المتاحة بالكامل على المسرعات الحديثة. ومن خلال التعميم إلى $K' > 1$ ، تقوم الخوارمة فعليًا بنقل الاختناق بعيدًا عن مرحلة الترتيب دون تحمل تكلفة إضافية كبيرة في مرحلة التصفية، طالما ظلت المرحلة الأولى مقيدة بالذاكرة.

يؤكد المؤلفون أن تحليلهم النظري يوفر أساسًا صارمًا لاختيار المعلمات، مما يوفر واجهة سهلة الاستخدام (approx_top_k(array, K, recall_target)) لا تتطلب أي ضبط يدوي. وبينما يركز التقييم التجريبي على TPUv5e، يؤكد المؤلفون أن الخوارزمية، ونمذجة الأداء، والتحليل عامة وقابلة للتطبيق على منصات أجهزة أخرى مثل GPUs. يوضح العمل أن التعميم الخوارزمي المنهجي، مقترنًا بالتنفيذ الواعي بالعتاد (الدمج/fusion)، يمكن أن يزيل الـ Top-K كعائق في تدريب واستنتاج النماذج واسعة النطاق.

الطريقة القديمة: مرشح "واحد في كل مرة"

الفكرة الجديدة: مرشح "Top-K"

"سحر" الأجهزة (Hardware)

النتائج: تسريع الذكاء الاصطناعي

أبحاث مشابهة