The quadratic complexity of self-attention limits its applicability and scalability on large unstructured meshes. We introduce Fast Low-rank Attention Routing Engine (FLARE), a linear complexity self-attention mechanism that routes attention through fixed-length latent sequences. Each attention head performs global communication among $N$ tokens by projecting the input sequence onto a fixed length latent sequence of $M \ll N$ tokens using learnable query tokens. By routing attention through a bottleneck sequence, FLARE learns a low-rank form of attention that can be applied at $O(NM)$ cost. FLARE not only scales to unprecedented problem sizes, but also delivers superior accuracy compared to state-of-the-art neural PDE surrogates across diverse benchmarks. We also release a new additive manufacturing dataset to spur further research. Our code is available at https://github.com/vpuri3/FLARE.py.
पारंपरिक स्व-अटेंशन तंत्र की द्विघात जटिलता बड़े पैमाने पर अनुरचित जाल पर इसकी प्रयोज्यता और स्केलेबिलिटी को सीमित करती है। यह पेपर फास्ट लो-रैंक अटेंशन राउटिंग इंजन (FLARE) प्रस्तुत करता है, जो निश्चित लंबाई के अव्यक्त अनुक्रम के माध्यम से अटेंशन को रूट करके रैखिक जटिलता का स्व-अटेंशन तंत्र है। प्रत्येक अटेंशन हेड सीखने योग्य क्वेरी टोकन का उपयोग करके इनपुट अनुक्रम को लंबाई M≪N के निश्चित लंबाई अव्यक्त अनुक्रम पर प्रक्षेपित करता है, जिससे N टोकन के बीच वैश्विक संचार प्राप्त होता है। बोतल गर्दन अनुक्रम राउटिंग अटेंशन के माध्यम से, FLARE कम-रैंक रूप का अटेंशन सीखता है जिसे O(NM) की लागत पर लागू किया जा सकता है। FLARE न केवल अभूतपूर्व समस्या आकार तक स्केल कर सकता है, बल्कि कई बेंचमार्क में अत्याधुनिक तंत्रिका PDE प्रॉक्सी मॉडल की तुलना में बेहतर सटीकता प्रदान करता है।
मुख्य समस्या: पारंपरिक Transformer के स्व-अटेंशन तंत्र में O(N²) समय और मेमोरी जटिलता है, जो बड़े पैमाने पर अनुरचित जाल (जैसे भौतिकी सिमुलेशन में बिंदु बादल और जाल) पर इसके अनुप्रयोग को गंभीरता से सीमित करती है।
अनुप्रयोग की महत्ता: आंशिक अवकल समीकरण (PDE) प्रॉक्सी मॉडलिंग में, 3D बिंदु बादल में प्रत्येक बिंदु को एक टोकन माना जाता है, जिसमें ज्यामितीय और भौतिक मात्राएं (जैसे निर्देशांक, सामान्य सदिश, सामग्री गुण) जैसी विशेषताएं होती हैं। उच्च निष्ठा भौतिकी प्रणाली सिमुलेशन की लागत अत्यधिक है, मशीन लर्निंग प्रॉक्सी मॉडल तेजी से सन्निकटन का विकल्प प्रदान करते हैं।
मौजूदा विधियों की सीमाएं:
PerceiverIO: केवल एकल एन्कोडिंग और डिकोडिंग करता है, संभावित बोतल गर्दन सटीकता को सीमित कर सकती है
Transolver: हेड में प्रक्षेपण भार साझा करता है, मौजूदा GPU कर्नेल के साथ स्केलिंग स्केल्ड डॉट प्रोडक्ट अटेंशन का उपयोग नहीं कर सकता
LNO: केवल एकल प्रक्षेपण लागू करता है, गहरे मॉडल क्षमता की कमी है
अनुसंधान प्रेरणा: एक ऐसा अटेंशन तंत्र विकसित करना जो वैश्विक संचार क्षमता बनाए रखे लेकिन रैखिक जटिलता हो, जिससे Transformer लाखों बिंदुओं वाली ज्यामिति को संभाल सके।
रैखिक जटिलता टोकन मिश्रण: FLARE स्व-अटेंशन तंत्र प्रस्तुत करता है जो कम-रैंक प्रक्षेपण और पुनर्निर्माण के माध्यम से पूर्ण स्व-अटेंशन को प्रतिस्थापित करके रैखिक जटिलता प्राप्त करता है।
उत्कृष्ट सटीकता: कई PDE बेंचमार्क में, FLARE कम पैरामीटर और कम कम्प्यूटेशनल जटिलता के साथ अग्रणी तंत्रिका प्रॉक्सी मॉडल से बेहतर भविष्यवाणी सटीकता प्राप्त करता है।
अभूतपूर्व स्केलेबिलिटी: FLARE पूरी तरह से मानक फ्यूज्ड अटेंशन प्राइमिटिव पर निर्मित है, उच्च GPU उपयोग सुनिश्चित करता है, लाखों बिंदु अनुरचित जाल के अंत-से-अंत प्रशिक्षण का समर्थन करता है।
नया बेंचमार्क डेटासेट: बड़े पैमाने पर उच्च-रिज़ॉल्यूशन धातु योज्य विनिर्माण डेटासेट जारी करता है, अवशिष्ट विस्थापन भविष्यवाणी अनुसंधान के लिए।
इनपुट अनुक्रम X ∈ R^(N×C) दिया गया है, जहां N टोकन की संख्या है, C विशेषता आयाम है, FLARE का लक्ष्य एक रैखिक जटिलता अटेंशन तंत्र सीखना है जो कुशल वैश्विक टोकन संचार प्राप्त करता है।
हेड-वार स्वतंत्र प्रक्षेपण: Transolver के साझा प्रक्षेपण भार के विपरीत, FLARE प्रत्येक हेड को अव्यक्त टोकन का एक अलग स्लाइस आवंटित करता है, जिससे प्रत्येक हेड स्वतंत्र अटेंशन संबंध सीख सकता है।
गहरा अवशिष्ट MLP: कुंजी/मूल्य प्रक्षेपण के लिए गहरे अवशिष्ट नेटवर्क का उपयोग करता है, सरल रैखिक परत की तुलना में उच्च-क्रम विशेषता इंटरैक्शन सीख सकता है।
सममित एन्कोड-डिकोड डिजाइन: एन्कोडिंग और डिकोडिंग संचालन की सममितता स्थिर सूचना प्रवाह को बढ़ावा देती है।
संगत फ्यूज्ड कर्नेल: पूरी तरह से मानक SDPA संचालन पर आधारित, Flash Attention जैसे अनुकूलन एल्गोरिदम का उपयोग कर सकता है।
FLARE एकल H100 GPU पर लाखों बिंदु DrivAerML डेटासेट को प्रशिक्षित करने में सफल रहा, यह पहला अटेंशन-आधारित तंत्रिका प्रॉक्सी मॉडल है जो मेमोरी अनलोडिंग या वितरित कंप्यूटिंग के बिना लाखों बिंदुओं को संभालता है।
पेपर Transformer, तंत्रिका ऑपरेटर, कुशल अटेंशन तंत्र आदि संबंधित क्षेत्रों के महत्वपूर्ण कार्यों का हवाला देता है, जो इस अनुसंधान के लिए ठोस सैद्धांतिक आधार और तुलना बेंचमार्क प्रदान करता है।
समग्र मूल्यांकन: यह एक उच्च-गुणवत्ता वाला अनुसंधान पेपर है जो Transformer स्केलेबिलिटी समस्या को हल करने के लिए एक नवाचारी समाधान प्रस्तुत करता है। FLARE विधि न केवल सैद्धांतिक रूप से कम-रैंक अपघटन की सुंदर व्याख्या रखती है, बल्कि व्यावहारिक रूप से उत्कृष्ट प्रदर्शन प्रदर्शित करती है। पेपर का प्रयोगात्मक डिजाइन व्यापक है, सैद्धांतिक विश्लेषण गहन है, और बड़े पैमाने पर ज्यामितीय गहन शिक्षा और वैज्ञानिक कंप्यूटिंग को आगे बढ़ाने में महत्वपूर्ण महत्व है।