Differentiable astrophysics at scale: solving and differentiating ODE ensembles on the GPU
تُقدم الورقة البحثية مكتبة GRADSOLVE، وهي مكتبة برمجية بلغة JAX تستفيد من التوازي في وحدات معالجة الرسومات لحل وتفاضل مجموعات من المعادلات التفاضلية العادية للنماذج الفلكية، محققةً تسريعاً يصل إلى 1,500 مرة مقارنة بالتطبيقات التسلسلية على وحدة المعالجة المركزية، ومتفوقةً بشكل كبير على حلول المعادلات التفاضلية العادية الحالية في JAX لتمكين التحليلات الروتينية القائمة على التدرج لملايين المسارات.
دخل علم الفلك الحديث عصرًا لا يكتفي فيه العلماء بمجرد النظر إلى السماء، بل يقومون بمحاكاتها بدقة مذهلة لفهم كيفية عمل الكون. وللتحقق من أفكارهم، يبني الباحثون نماذج رياضية معقدة تصف كل شيء، بدءًا من المسارات الدوامية للنجوم في مجرتنا وصولاً إلى توسع الفضاء نفسه. وتعتمد هذه النماذج غالبًا على حل نوع محدد من الألغاز الرياضية يسمى "المعادلة التفاضلية العادية". فكر في هذه المعادلات كأنها مجموعة من التعليمات التي تخبر الحاسوب كيف يتغير النظام بمرور الوقت، مثل كيفية تحرك نجم تحت تأثير الجاذبية أو كيف يتمدد الكون مع تقدم عمره. ولإيجاد أفضل نموذج ممكن، يجب على علماء الفلك تشغيل هذه التعليمات ملايين المرات، مع تعديل الأرقام قليلاً في كل مرة لمعرفة أي نسخة تطابق الواقع بشكل وثيق. هذه العملية هي العمود الفقري للاكتشافات الحديثة، لكنها اصطدمت بحائط مسدود: فالعدد الهائل من الحسابات المطلوبة ضخم جدًا لدرجة أن حتى أسرع الحواسيب الفائقة تجد صعوبة في إنهاء المهمة في وقت معقول.
لقد نجح أحد الباحثين الآن في كسر هذا الاختناق عبر نقل هذه الحسابات من معالجات الكمبيوتر القياسية إلى وحدات معالجة الرسومات، وهي الرقائق القوية التي صُممت في الأصل لتشغيل ألعاب الفيديو. فقد طور أداة برمجية جديدة تسمى gradsolve تعامل كل عملية محاكاة من ملايين عمليات المحاكاة كعامل منفصل، حيث تخصص لكل منها مسارًا مخصصًا على شريحة الكمبيوتر. في الطرق التقليدية، يحاول الحاسوب حل هذه المشكلات واحدة تلو الأخرى أو في مجموعات جامدة، مما يجبر عمليات المحاكاة السريعة غالبًا على الانتظار خاملة حتى تلحق بها العمليات الأبطأ. أما هذا النهج الجديد فيسمح لكل عملية محاكاة بالعمل بسرعتها الخاصة، بحيث تنتهي تمامًا عندما تكون جاهزة دون تعطيل أي شخص آخر. والنتيجة هي قفزة هائلة في السرعة تحول ما كان يستغرق أيامًا في السابق إلى مهمة يمكن إنجازها في ثوانٍ.
اختبر الباحث هذه الطريقة على ثلاث مشكلات مختلفة تمامًا من مختلف مجالات علم الفلك لإثبات تعدد استخداماتها. أولاً، درس مسارات النجوم التي تُنتزع من عنقود كروي لتشكل تيارًا طويلًا ونحيفًا عبر درب التبانة. ومن خلال محاكاة حركة مليون نجم في آن واحد، كانت أداة البحث الجديدة أسرع بـ 1500 مرة من أفضل البرامج الموجودة التي تعمل على كمبيوتر قياسي يحتوي على 128 معالجًا. ثانيًا، طبقها على تاريخ توسع الكون، وتحديدًا لاختبار نظريات حول الطاقة المظلمة، تلك القوة الغامضة التي تدفع الكون للتباعد. وهنا، قامت الأداة بحساب المسافات لمليون نسخة مختلفة من الكون، وتفوقت مرة أخرى على الطرق القديمة بمعامل قدره 1500 مقارنة بنواة حاسوب واحدة. وأخيرًا، فحصت حركة الالتواء، أو المبادرة (precession)، لثقبين أسودين يلتفان حول بعضهما البعض قبل الاصطدام. وبالنسبة لمجموعة مكونة من مليون ثقب أسود ثنائي، كانت البرمجية الجديدة أسرع بـ 760 مرة من الكود القياسي الذي يستخدمه علماء الفلك اليوم.
إن ما يجعل هذا الإنجاز مهمًا بشكل خاص هو أن الأداة الجديدة لا تعمل بسرعة أكبر فحسب، بل تحسب أيضًا المشتقات اللازمة، أو معدلات التغيير، في نفس المرحلة. في الماضي، كان العثور على هذه المعدلات يتطلب عملية منفصلة ومكلفة بنفس القدر، مما كان سيضاعف الوقت المطلوب. ومن خلال التعامل مع المحاكاة وحساسيتها للتغيرات في آن واحد، أزال الباحث عائقًا رئيسيًا أمام التحليل الإحصائي المعقد. إن مكاسب السرعة كبيرة جدًا لدرجة أن التحليلات التي كانت في السابق مكلفة حاسوبيًا للغاية بحيث يصعب القيام بها، أصبحت الآن روتينية. فعلى سبيل المثال، المهمة التي قد تستغرق من نواة حاسوب واحدة قرابة العام لإكمالها، يمكن الآن إنهاؤها على بطاقة رسومات واحدة في غضون ثوانٍ.
لقد حرص الباحث على ضمان أن هذه الزيادة الهائلة في السرعة لا تأتي على حساب الدقة. فقد وضع معايير صارمة لمدى قرب النتائج من حل مرجعي معروف وعالي الدقة. وفي كل اختبار، استوفت البرمجية الجديدة هذه المتطلبات الصارمة للدقة، مما أثبت أن السرعة جاءت من طريقة أذكى لتنظيم العمل وليس من اختصار الطرق. وتظهر الدراسة بوضوح أنه بالنسبة للمشكلات التي تتضمن ملايين الحسابات المستقلة التي لا تتطلب خطوات رياضية معقدة وصعبة (stiff)، فإن هذا النهج الجديد يتفوق على الأساليب الحالية المتطورة. وهذا يثبت أنه من خلال إعادة التفكير في كيفية تعامل الحواسيب مع مجموعات ضخمة من عمليات المحاكاة، يمكن لعلماء الفلك الآن استكشاف الكون بمستوى من التفصيل والسرعة كان بعيد المنال سابقًا، مما يفتح الباب أمام نماذج أكثر دقة لكل شيء، من بنية مجرتنا إلى الطبيعة الأساسية للطاقة المظلمة.
ملخص تقني: الفيزياء الفلكية القابلة للتفاضل على نطاق واسع
بيان المشكلة تعتمد عمليات الاستدلال في الفيزياء الفلكية الحديثة بشكل متزايد على الطرق القائمة على التدرج (gradient-based methods)، مثل مونت كارلو هاملتونيان (HMC) وتوقعات فيشر (Fisher forecasts)، والتي تتطلب مشتقات تنبؤات النموذج بالنسبة لمعلماته. وتتمثل إحدى العقبات الحسابية الشائعة في هذه التحليلات في الحاجة إلى حل مجموعات من المعادلات التفاضلية العادية (ODEs)—التي تتراوح من آلاف إلى ملايين المسارات—حيث يتوافق كل مسار مع مجموعة مختلفة من المعلمات أو الشروط الأولية.
تواجه النهج الحالية مشكلتين أساسيتين:
التنفيذ المتسلسل على وحدة المعالجة المركزية (CPU): إن حل هذه المجموعات على وحدات المعالجة المركزية عملية مكلفة حاسوبياً، مما يحد غالباً من نطاق التحليلات.
قيود التوجيه الشعاعي (Vectorization) على وحدة معالجة الرسومات (GPU): تعمل مكتبات المعادلات التفاضلية العادية القابلة للتفاضل في JAX، مثل diffrax المتطورة، على توجيه المجموعة كصفوف (array) واحدة. وفي الحلول التكيفية (adaptive solvers)، حيث تختلف أحجام الخطوات بناءً على تقديرات الخطأ المحلية، يفرض هذا التنفيذ "المتزامن خطوة بخوة" (lock-step) على جميع المسارات انتظار أبطأ عضو في المجموعة لإنهاء خطوته. ونتيجة لذلك، تصبح التكلفة الحسابية محكومة بالمسار الذي يتطلب أكبر عدد من الخطوات بدلاً من المتوسط، مما يؤدي إلى عدم كفاءة كبيرة في وحدات معالجة الرسومات.
المنهجية يقدم المؤلف مكتبة gradsolve، وهي مكتبة مبنية على JAX مصمم لنقل تكامل مجموعات المعادلات التفاضلية العادية إلى وحدات معالجة الرسومات باستخدام بنية "مسار واحد لكل خيط" (one trajectory per thread).
التكامل لكل خيط (Per-Thread Integration): على عكس النهج الموجه شعاعياً، تقوم gradsolve بإنشاء نواة (kernel) لوحدة معالجة الرسومات حيث يقوم كل خيط بتكامل مسار واحد بشكل مستقل. يحافظ كل خيط على حالته الخاصة، وتقييماته الوسيطة، وحساسياته في ذاكرة خاصة. يسمح هذا للمسارات بالمضي قدماً بأحجام خطواتها التكيفية الخاصة دون انتظار الآخرين، مما يحصر التزامن فقط في مجموعات من 32 خيطاً (warps) على وحدات NVIDIA GPU.
الحساسيات الأمامية (Forward Sensitivities): تقوم المكتبة بتكامل معادلات الحساسية الأمامية (مشتقات الحالة بالنسبة للمعلمات) جنباً إلى جنب مع المعادلات التفاضلية العادية داخل نفس النواة. يعد هذا التفاضل في الوضع الأمامي (forward-mode differentiation) فعالاً للمشكلات ذات عدد قليل من المعلمات والعديد من المخرجات، مما يتجنب العبء الإضافي للذاكرة المطلوب في تفاضل الوضع العكسي (adjoint differentiation) الذي يتطلب تخزين أو إعادة حساب المسار الأمامي.
المخرج الكثيف (Dense Output): يوفر الحل مخرجاً كثيفاً، مما يسمح باستكمال (interpolation) الحل عند أوقات رصد تعسفية دون إجبار الحل على إنهاء خطواته عند تلك الأوقات المحددة.
بروتوكول المقارنة: لضمان عدالة اختبار الأداء، يقارن المؤلف gradsolve مقابل diffrax والأكواد المتسلسلة على وحدة المعالجة المركزية (مثل gala و scipy و precession) تحت متطلبات دقة صارمة ومحددة مسبقاً. تُقبل الأكواد فقط إذا استوفت المئوية الـ 99 من خطئها عبر المجموعة عتبة الدقة الفيزيائية (على سبيل المثال، الخطأ النسبي في المدار أو المسافة).
المساهمات الرئيسية
تطوير المكتبة: إصدار gradsolve كـ مكتبة أصلية لـ JAX تولد نوى GPU لكل خيط لتكامل المعادلات التفاضلية العادية التكيفي وحساب الحساسية الأمامية.
التحول الهيكلي: إثبات أن نهج "الخيط الواحد لكل مسار" يتفوق بشكل كبير على نهج المجموعات الموجهة شعاعياً على وحدات معالجة الرسومات بالنسبة لمجموعات المعادلات التفاضلية غير المتصلبة (non-stiff) وذات الأنظمة الصغيرة.
اختبار أداء شامل: تقديم مقارنات أداء صارمة عبر ثلاث مجالات فيزيائية فلكية متميزة:
الطاقة المظلمة (Dark Energy): حساب تاريخ التوسع والمسافات لنماذج "الكمنتيسنس" (quintessence) المتغيرة، بما في ذلك حل قيود التسطح المكاني.
الثقوب السوداء الثنائية (Binary Black Holes): محاكاة بدارية الدوران (spin precession) عبر مجموعة من الثقوب السوداء الثنائية.
النتائج يسجل البحث عوامل تسريع لمجموعات مكونة من 106 مسار، مقارنة بـ gradsolve مقابل كل من أكواد وحدة المعالجة المركزية المتسلسلة (التي تعمل على 128 نواة من AMD EPYC 9745) وأسرع تكوين مقبول لـ diffrax على نفس وحدة NVIDIA H200 GPU.
مقابل وحدة المعالوة المركزية المتسلسلة:gradsolveأسرع بـ 8.5 إلى 1,500 مرة من كود وحدة المعالجة المركزية المتسلسل الذي يعمل على 128 نواة. مقابل نواة واحدة من وحدة المعالجة المركزية، يصل التسريع إلى عدة مراتب عشرية.
المدارات النجمية: أسرع بـ 8.5 مرة من 128 نواة.
الطاقة المظلمة: أسرع بـ 1,500 مرة من 128 نواة.
بدارية الدوران: أسرع بـ 760 مرة من 128 نواة.
مقابل diffrax (على نفس وحدة معالجة الرسومات):gradsolveأسرع بـ 11 إلى 15 مرة من diffrax للمجموعات الكبيرة (106 عضو) مع الحفاظ على دقة متطابقة.
المدارات النجمية: أسرع بـ 15 مرة (المدارات + المشتقات).
الطاقة المظلمة: أسرع بـ 11 مرة (بما في ذلك حل التسطح المكاني).
بدارية الدوران: أسرع بـ 11 مرة.
الدقة: استوفت جميع الأكواد المختبرة متطلبات الدقة المحددة مسبقاً. ظلت المئوية الـ 99 من الأخطاء لـ gradsolve ضمن الحدود الفيزيائية المحددة لكل تطبيق (على سبيل المثال، 10−4 للمواقع المدارية، 10−7 راديان لزوايا البدارية).
مصدر التسريع: عند مقارنة gradsolve و diffrax باستخدام نفس الطريقة العددية (مثل Verner 7 أو Tsit5) والتحمل، تكون محاولات الخطوات متطابقة تقريباً. يُعزى فرق الأداء إلى التنفيذ: تتجنب gradsolve عبء عرض نطاق الذاكرة (memory bandwidth) الناتج عن قراءة/كتابة حالة المجموعة بأكملها في الذاكرة العالمية عند كل خطوة، وتلغي عقوبة "الانتظار" الناتجة عن حلقة التنفيذ المتزامنة.
الأهمية والنطاق يزعم البحث أن gradsolve يجعل التحليلات القائمة على التدرج التي تتطلب مجموعات من ملايين التكاملات للمعادلات التفاضلية العادية أمراً "روتينياً" من خلال تقليل أوقات الحوسبة من ساعات إلى ثوانٍ على وحدة معالجة رسومات واحدة.
القابلية للتطبيق: هذا النهج مُحسّن للأنظمة غير المتصلبة (non-stiff) ذات متجهات الحالة الصغيرة (3–9 مكونات) وعدد قليل من المعلمات المتفاضلة (كفاءة الوضع الأمامي).
القيود: يشير البحث إلى أنه بالنسبة للمشكلات ذات متجهات الحالة الكبيرة (مثل محاكاة الهيدروديناميكا الشبكية أو محاكاة الجسيمات-الشبكة) أو الأنظمة المتصلبة (stiff systems) (مثل معادلات أينشتاين-بولتزمان)، فإن النهج الأخرى (مثل الحلول المصفوفية أو طرق Rosenbrock الضمنية) تكون أكثر ملاءمة. وبينما تتضمن gradsolve نواة Rosenbrock للمشكلات المتصلبة، فإن الاختبارات المرجعية الأساسية تركز على التطبيقات الفلكية غير المتصلبة.
الأثر على الاستدلال: في مثال احتمالية التوزيع للطاقة المظلمة، يقلل التسريع من تكلفة حل المعادلات التفاضلية العادية إلى جزء ضئيل من تكلفة التكرار الإجمالية، مما ينقل الاختناق الحسابي إلى حساب الاحتمالية (likelihood calculation) نفسه. وهذا يتيح استكشاف نماذج أكثر تعقيداً وفضاءات بارامترية أكبر ضمن أطر زمنية ممكنة.
الكود متاح للجمهور، مما يسمح للمجتمع الفلكي بدمج هذه الحلول عالية الأداء والقابلة للتفاضل في مسارات الاستدلال القائمة على JAX.