ملخص تقني: OmniMol

بيان المشكلة

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

يكمن التحدي الرئيسي الذي يتم معالجته في تطوير إمكانات ذرية متعلمة آلياً (MLIPs) فعالة. الطرق التقليدية مثل نظرية وظيفية الكثافة (DFT) مكلفة حاسوبياً، مما يحد من عمليات محاكاة الديناميكا الجزيئية (MD) واسعة النطاق وطويلة الأمد. تهدف الـ MLIPs إلى تقريب أسطح طاقة الوضع والقوى بتكلفة أقل بكثير. تفترض الورقة أن نموذجاً تأسيسياً (Foundation Model) تم تدريبه مسبقاً على سحب النقاط في فيزياء الجسيمات (تحديداً نفاثات الجسيمات - particle jets) يمكن نقله إلى الديناميكا الجزيئية، مما قد يسرع عملية التحسين ويحسن الدقة في الأنظمة ذات البيانات المحدودة.

المنهجية

البنية: محول النقطة والحافة (Point-Edge Transformer - PET)

بُني OmniMol عبر تكييف Omnilearned، وهو نموذج تأسيسي صُمم أصلاً لتصنيف وتوليد نفاثات الجسيمات في فيزياء الطاقات العالية (HEP). البنية الأساسية هي محول النقطة والحافة (PET)، الذي يربط بين الانتباه المحلي عبر أقرب $k$ جيران مع كتل المحول العالمية (all-to-all transformer blocks).

المكونات الهيكلية الرئيسية تشمل:

1. تضمينات المدخلات (Input Embeddings): يتم تضمين الذرات في فضاء الرموز (token space) الذي يجمع بين المعلومات الموضعية ($\vec{r}$)، والأعداد الذرية المنفصلة ($Z$)، وميزات إضافية (الشحنة، اللف المغزلي).
2. كتلة الانتباه المحلي (Local Attention Block): لكل ذرة، يتم إنشاء جوار محلي باستخدام أقرب $K$ جيران ($K=15$ للجزيئات، مقارنة بـ $K=10$ للنفاثات). تُحسب الميزات الفيزيائية الزوجية، بما في ذلك حدود المسافة، والقوى العكسية للمسافة، والدوال المتعلمة من التضمينات الذرية. تتم معالجة هذه الميزات بواسطة محول محلي صغير لإنشاء متجه تضمين محلي.
3. الانتباه العالمي مع انحياز التفاعل (Global Attention with Interaction Bias): يتضمن آلية الانتباه الذاتي العالمي انحيازاً صريحاً مشتقاً من الميزات الفيزيائية الزوجية. يتم تعديل لوغاريتمات الانتباه كـ $A^*_{ij} = A_{ij} + B_{ij}$، حيث $B_{ij}$ هو حد انحياز مُضمن في MLP. هذا "انحياز مصفوفة التفاعل في الانتباه" يضخ المبادئ الفيزيائية الزوجية مباشرة في المحول، مما يوجه الشبكة نحو الجوارات ذات المعنى الفيزيائي دون التضحية بالقدرة التعبيرية.
4. رؤوس المخرجات (Output Heads): تمت إعادة توظيف الرأس التوليدي لـ Omnilearned لمهمتين:
   • التنبؤ بالقوة: رأس متوافق مع التبديل (permutation-equivariant) يتنبأ بالقوى لكل ذرة.
   • التنبؤ بالطاقة: رأس يتنبأ بتصحيحات الطاقة لكل ذرة، والتي تُجمع لتعطي إجمالي الطاقة الجزيئية، مما يحافظ على المبادئ التراكمية (extensive priors).

قيود الثبات والتماثل (Invariance and Conservation Constraints)

لتحقيق القيود الفيزيائية، عالج المؤلفون مطلبين:

• حفظ الطاقة: لا يتم التنبؤ بالقوى مباشرة، بل تُحسب عبر الانتشار العكسي لمخرج الطاقة ($\vec{F}_i = \nabla_{\vec{r}_i} E$). يضمن هذا حفظ الطاقة بدقة ولكنه يزيد التكلفة الحاسوبية أثناء التدريب (يتطلب انتشاراً عكسياً مزدوجاً). بناءً على ذلك، يُطبق هذا القيد فقط على نسخة النموذج "الصغير".
• التماثل الدوراني (Rotational Equivariance): البنية القياسية ليست متماثلة بطبيعتها لأن فروق الإحداثيات الخام تُغذى في الـ MLPs. لمعالجة ذلك، قدم المؤلفون نسخة "متماثلة ومحافظة". تقوم هذه النسخة بإزالة حدود فروق الإحداثيات المباشرة من الميزات الزوجية، وبدلاً من ذلك تدمج المعلومات الزاوية (جيب تمام الزوايا المتشكلة من المتجهات بين الذرات المجاورة) في الكتلة المحلية. هذا التعديل يحافظ على التماثل مع استعادة الأداء المرتبط بإزالة حدود الإحداثيات بشكل كبير.

استراتيجيات التدريب والضبط الدقيق

تم ضبط النموذج دقيقاً على مجموعة بيانات oMol (تحديداً المجموعات الفرعية oMol-25، oMol-4M، oMol-100M، وoMol-140M). تم استكشاف استراتيجيتين:

1. LoRA (التكيف منخفض الرتبة): يتم تجميد أوزان متن (backbone) PET المدرب مسبقاً. تُدخل محولات منخفضة الرتبة فقط لمصفوفات جسم المحول ($W_Q, W_K, W_V, W_O, W_{MLP}$)، جنباً إلى جنب مع تدريب مشفرات المدخلات الجزيئية، وMLP الانحياز، ورؤوس المهام. كما تُضاف طبقة "تكييف التضمين" لتعديل التضمينات المتعلمة.
2. الضبط الدقيق الكامل (Full Fine-Tuning): يتم إلغاء تجميد جميع الأوزان في الجسم ومشفرات المدخلات، بينما يتم تدريب رؤوس المهام من الصفر.

يهدف هدف التدريب إلى تقليل مجموع متوسط الأخطاء المطلقة (MAE) للطاقات والقوى، مع إعطاء وزن أكبر للقوى ($\lambda_F = 10$).

النتائج الرئيسية

الأداء على oMol

• التدريب الكامل: عندما يتم التدريب على مجموعات بيانات كبيرة (oMol-4M وoMol-100M/140M)، يحقق OmniMol أداءً تنافسياً مع أفضل الـ MLIPs. على سبيل المثال، في oMol-140M، يحقق نموذج OmniMol-large متوسط خطأ طاقة قدره 1.04 meV/atom ومتوسط خطأ قوة قدره 13.59 meV/Å.
• نظام البيانات المحدودة: تُلاحظ أكبر المكاسب عند محدودية بيانات التدريب. عندما يتم الضبط الدقيق على 100 ألف جزيء أو مع عدد قليل جداً من الدورات (دورتين تمرير) فوق oMol-4M، تتفوق متغيرات Omni_Mol المدربة مسبقاً بشكل كبير على النماذج المدربة من الصفر.
  • في مجموعة فرعية مكونة من 100 ألف، أدى التدريب المسبق إلى تحسين متوسط خطأ الطاقة بنسبة تصل إلى 29.4% ومتوسط خطأ القوة بنسبة 26.9% للنموذج المتوسط.
  • مع دورتين فقط على oMol-4M، أظهر النموذج المتوسط تحسناً بنسبة 54.6% في متوسط خطأ الطاقة و 56.9% في متوسط خطأ القوة مقارنة بنظيره غير المدرب مسبقاً.
• النسخة المتماثلة/المحافظة: تظهر النسخة المتماثلة والمحافظة أداءً محسناً بشكل ملحوظ (خاصة بالنسبة للقوى) في أنظمة البيانات المحدودة، رغم أن هذه الميزة تتضاءل مع زيادة حجم مجموعة البيانات.

التوسع وسرعة الاستدلال

• التوسع (Scaling): يتبع OmniMol قانون القوة (power-law scaling) الواضح مع حجم النموذج، حيث لا تظهر أي علامات على التشبع حتى مليار معلمة، وهو ما يتوافق مع النتائج الأخيرة على المحولات المستخدمة في MLIPs.
• سرعة الاستدلال: رغم عدد المعلمات الكبير، يُظهر OmniMol سرعات استدلال سريعة بشكل فريد بسبب التحسينات العتادية للمحولات. على وحدة GPU من نوع A100 للأنظمة التي تحتوي على ~100 ذرة، يعتبر OmniMol-medium أسرع بحوالي 3 مرات من خطوط الأساس المقارنة من الشبكات العصبية الرسومية (GNNs) (مثل end-to-end equivariant GNNs و AllScAIP-md) مع الحفاظ على دقة تنافسية (بفارق خطأ طاقة قدره ~0.7 meV/atom فقط عن AllScAIP-md).

الأهمية والادعاءات

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

1. النقل عبر المجالات قابل للتطبيق: يمكن لنموذج تأسيسي بُني لفيزياء الجسيمات أن ينتقل بفعالية إلى الكيمياء الجزيئية، مما يشير إلى أن بنيات سحب النقاط الأساسية تشترك في ميزات قابلة للتعلم عبر مقاييس فيزيائية مختلفة تماماً.
2. الانحياز الاستقرائي يسرع التعلم: يعمل التدريب المسبق كـ "انحياز استقرائي" قوي. تماماً كما يساعد التماثل (equivariance) عند ندرة البيانات، فإن "الدرس المرير" (bitter lesson) للتدريب المسبق يسمح بتحسين سريع ودقة أعلى عندما تكون بيانات التدريب محدودة.
3. الكفاءة: يتيح النقل الهيكلي سرعات استدلال سريعة بشكل فريد، وهو أمر بالغ الأهمية للتطبيقات التي تتطلب استكشافاً سريعاً لمساحات التصميم، مثل اكتشاف الأدوية للجزيئات الصغيرة.

يخلص المؤلفون إلى أنه بينما تركز الدراسة على MLIPs، فإن الدروس المتعلقة بنماذج التأسيس لسحب النقاط قد يكون لها فائدة واسعة النطاق عبر المجالات العلمية حيث توصف الأنظمة كمجموعات غير مرتبة من الأجسام المتفاعلة. هم لا يدعون التفوق العالمي على جميع الطرق الموجودة في جميع الأنظمة، بل يسلطون الضوء على مزايا محددة في سيناريوهات البيانات المحدودة وسرعة الاستدلال.