2025-11-17T15:49:13.397134

FLARE: Fast Low-rank Attention Routing Engine

Puri, Joglekar, Ferguson et al.
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.
academic

FLARE: फास्ट लो-रैंक अटेंशन राउटिंग इंजन

बुनियादी जानकारी

  • पेपर ID: 2508.12594
  • शीर्षक: FLARE: Fast Low-rank Attention Routing Engine
  • लेखक: Vedant Puri, Aditya Joglekar, Kevin Ferguson, Yu-hsuan Chen, Yongjie Jessica Zhang, Levent Burak Kara (Carnegie Mellon University)
  • वर्गीकरण: cs.LG (मशीन लर्निंग)
  • प्रकाशन समय: 15 अक्टूबर 2025 (arXiv v2)
  • पेपर लिंक: https://arxiv.org/abs/2508.12594

सारांश

पारंपरिक स्व-अटेंशन तंत्र की द्विघात जटिलता बड़े पैमाने पर अनुरचित जाल पर इसकी प्रयोज्यता और स्केलेबिलिटी को सीमित करती है। यह पेपर फास्ट लो-रैंक अटेंशन राउटिंग इंजन (FLARE) प्रस्तुत करता है, जो निश्चित लंबाई के अव्यक्त अनुक्रम के माध्यम से अटेंशन को रूट करके रैखिक जटिलता का स्व-अटेंशन तंत्र है। प्रत्येक अटेंशन हेड सीखने योग्य क्वेरी टोकन का उपयोग करके इनपुट अनुक्रम को लंबाई M≪N के निश्चित लंबाई अव्यक्त अनुक्रम पर प्रक्षेपित करता है, जिससे N टोकन के बीच वैश्विक संचार प्राप्त होता है। बोतल गर्दन अनुक्रम राउटिंग अटेंशन के माध्यम से, FLARE कम-रैंक रूप का अटेंशन सीखता है जिसे O(NM) की लागत पर लागू किया जा सकता है। FLARE न केवल अभूतपूर्व समस्या आकार तक स्केल कर सकता है, बल्कि कई बेंचमार्क में अत्याधुनिक तंत्रिका PDE प्रॉक्सी मॉडल की तुलना में बेहतर सटीकता प्रदान करता है।

अनुसंधान पृष्ठभूमि और प्रेरणा

समस्या की पृष्ठभूमि

  1. मुख्य समस्या: पारंपरिक Transformer के स्व-अटेंशन तंत्र में O(N²) समय और मेमोरी जटिलता है, जो बड़े पैमाने पर अनुरचित जाल (जैसे भौतिकी सिमुलेशन में बिंदु बादल और जाल) पर इसके अनुप्रयोग को गंभीरता से सीमित करती है।
  2. अनुप्रयोग की महत्ता: आंशिक अवकल समीकरण (PDE) प्रॉक्सी मॉडलिंग में, 3D बिंदु बादल में प्रत्येक बिंदु को एक टोकन माना जाता है, जिसमें ज्यामितीय और भौतिक मात्राएं (जैसे निर्देशांक, सामान्य सदिश, सामग्री गुण) जैसी विशेषताएं होती हैं। उच्च निष्ठा भौतिकी प्रणाली सिमुलेशन की लागत अत्यधिक है, मशीन लर्निंग प्रॉक्सी मॉडल तेजी से सन्निकटन का विकल्प प्रदान करते हैं।
  3. मौजूदा विधियों की सीमाएं:
    • PerceiverIO: केवल एकल एन्कोडिंग और डिकोडिंग करता है, संभावित बोतल गर्दन सटीकता को सीमित कर सकती है
    • Transolver: हेड में प्रक्षेपण भार साझा करता है, मौजूदा GPU कर्नेल के साथ स्केलिंग स्केल्ड डॉट प्रोडक्ट अटेंशन का उपयोग नहीं कर सकता
    • LNO: केवल एकल प्रक्षेपण लागू करता है, गहरे मॉडल क्षमता की कमी है
  4. अनुसंधान प्रेरणा: एक ऐसा अटेंशन तंत्र विकसित करना जो वैश्विक संचार क्षमता बनाए रखे लेकिन रैखिक जटिलता हो, जिससे Transformer लाखों बिंदुओं वाली ज्यामिति को संभाल सके।

मुख्य योगदान

  1. रैखिक जटिलता टोकन मिश्रण: FLARE स्व-अटेंशन तंत्र प्रस्तुत करता है जो कम-रैंक प्रक्षेपण और पुनर्निर्माण के माध्यम से पूर्ण स्व-अटेंशन को प्रतिस्थापित करके रैखिक जटिलता प्राप्त करता है।
  2. उत्कृष्ट सटीकता: कई PDE बेंचमार्क में, FLARE कम पैरामीटर और कम कम्प्यूटेशनल जटिलता के साथ अग्रणी तंत्रिका प्रॉक्सी मॉडल से बेहतर भविष्यवाणी सटीकता प्राप्त करता है।
  3. अभूतपूर्व स्केलेबिलिटी: FLARE पूरी तरह से मानक फ्यूज्ड अटेंशन प्राइमिटिव पर निर्मित है, उच्च GPU उपयोग सुनिश्चित करता है, लाखों बिंदु अनुरचित जाल के अंत-से-अंत प्रशिक्षण का समर्थन करता है।
  4. नया बेंचमार्क डेटासेट: बड़े पैमाने पर उच्च-रिज़ॉल्यूशन धातु योज्य विनिर्माण डेटासेट जारी करता है, अवशिष्ट विस्थापन भविष्यवाणी अनुसंधान के लिए।

विधि विवरण

कार्य परिभाषा

इनपुट अनुक्रम X ∈ R^(N×C) दिया गया है, जहां N टोकन की संख्या है, C विशेषता आयाम है, FLARE का लक्ष्य एक रैखिक जटिलता अटेंशन तंत्र सीखना है जो कुशल वैश्विक टोकन संचार प्राप्त करता है।

मॉडल आर्किटेक्चर

FLARE मुख्य तंत्र

FLARE M≪N सीखने योग्य अव्यक्त टोकन प्रस्तुत करता है जो सूचना विनिमय के लिए बोतल गर्दन के रूप में कार्य करते हैं, दो चरणों को शामिल करते हैं:

  1. एन्कोडिंग चरण: इनपुट अनुक्रम क्रॉस-अटेंशन के माध्यम से अव्यक्त टोकन में प्रक्षेपित होता है
    Z_h = SDPA(Q_h, K_h, V_h, s=1)
    

    जहां Q_h ∈ R^(M×D) सीखने योग्य क्वेरी मैट्रिक्स है, K_h, V_h ∈ R^(N×D)
  2. डिकोडिंग चरण: अव्यक्त टोकन इनपुट अनुक्रम में वापस प्रक्षेपित होते हैं
    Y_h = SDPA(K_h, Q_h, Z_h, s=1)
    

कम-रैंक संचार मैट्रिक्स

पूरी प्रक्रिया समतुल्य है:

Y_h = (W_decode,h · W_encode,h) · V_h

जहां:

  • W_encode,h = softmax(Q_h · K_h^T) ∈ R^(M×N)
  • W_decode,h = softmax(K_h · Q_h^T) ∈ R^(N×M)
  • W_h = W_decode,h · W_encode,h ∈ R^(N×N) अधिकतम M रैंक का वैश्विक संचार मैट्रिक्स है

FLARE ब्लॉक संरचना

X = X + FLARE(LayerNorm(X))
X = X + ResMLP(LayerNorm(X))

तकनीकी नवाचार

  1. हेड-वार स्वतंत्र प्रक्षेपण: Transolver के साझा प्रक्षेपण भार के विपरीत, FLARE प्रत्येक हेड को अव्यक्त टोकन का एक अलग स्लाइस आवंटित करता है, जिससे प्रत्येक हेड स्वतंत्र अटेंशन संबंध सीख सकता है।
  2. गहरा अवशिष्ट MLP: कुंजी/मूल्य प्रक्षेपण के लिए गहरे अवशिष्ट नेटवर्क का उपयोग करता है, सरल रैखिक परत की तुलना में उच्च-क्रम विशेषता इंटरैक्शन सीख सकता है।
  3. सममित एन्कोड-डिकोड डिजाइन: एन्कोडिंग और डिकोडिंग संचालन की सममितता स्थिर सूचना प्रवाह को बढ़ावा देती है।
  4. संगत फ्यूज्ड कर्नेल: पूरी तरह से मानक SDPA संचालन पर आधारित, Flash Attention जैसे अनुकूलन एल्गोरिदम का उपयोग कर सकता है।

प्रयोगात्मक सेटअप

डेटासेट

पेपर 6 बेंचमार्क डेटासेट और 1 नए प्रस्तावित डेटासेट का मूल्यांकन करता है:

डेटासेटआयामजाल प्रकारबिंदु संख्याइनपुट/आउटपुट विशेषताएंप्रशिक्षण/परीक्षण नमूने
Elasticity2Dअनुरचित9722/11000/200
Darcy2Dसंरचित7,2252/11000/200
Airfoil2Dसंरचित11,2712/11000/200
Pipe2Dसंरचित16,6412/11000/200
DrivAerML-40k3Dअनुरचित40,0003/1387/97
LPBF3Dअनुरचित1,000-50,0003/11100/290

मूल्यांकन मेट्रिक्स

मुख्य रूप से सापेक्ष L2 त्रुटि का उपयोग करता है:

Relative L2 = ||û - u||₂ / ||u||₂

तुलना विधियां

  • सामान्य अटेंशन मॉडल: Vanilla Transformer, PerceiverIO
  • अटेंशन-आधारित PDE प्रॉक्सी: Transolver, LNO
  • तंत्रिका ऑपरेटर: GNOT

कार्यान्वयन विवरण

  • अनुकूलक: AdamW (β₁=0.9, β₂=0.999)
  • सीखने की दर अनुसूची: OneCycleLR, शिखर सीखने की दर 10⁻³
  • प्रशिक्षण युग: 2D समस्याओं के लिए 500, LPBF के लिए 250
  • बैच आकार: 2D समस्याओं के लिए 2, 3D समस्याओं के लिए 1

प्रयोगात्मक परिणाम

मुख्य परिणाम

FLARE सभी बेंचमार्क में सर्वोत्तम या द्वितीय-सर्वोत्तम परिणाम प्राप्त करता है:

मॉडलElasticityDarcyAirfoilPipeDrivAerML-40kLPBF
Vanilla Transformer5.374.386.28
PerceiverIO23.421.51627.1476056.3
GNOT13.316.91035.8911524.3
LNO9.257.6417.88.1014624.7
Transolver बिना conv6.4018.68.244.8770.520.4
Transolver conv के साथ\5.945.503.90\\
FLARE (हमारा)3.385.104.282.8560.818.5

नोट: संख्याएं सापेक्ष L2 त्रुटि हैं (×10⁻³)

लाखों बिंदु ज्यामिति प्रयोग

FLARE एकल H100 GPU पर लाखों बिंदु DrivAerML डेटासेट को प्रशिक्षित करने में सफल रहा, यह पहला अटेंशन-आधारित तंत्रिका प्रॉक्सी मॉडल है जो मेमोरी अनलोडिंग या वितरित कंप्यूटिंग के बिना लाखों बिंदुओं को संभालता है।

विलोपन प्रयोग

  1. ब्लॉक संख्या (B) और अव्यक्त टोकन संख्या (M) का प्रभाव:
    • ब्लॉक संख्या में वृद्धि सापेक्ष त्रुटि को लगातार कम करती है
    • M में वृद्धि आमतौर पर प्रदर्शन में सुधार करती है, लेकिन प्रवृत्ति कड़ाई से एकरस नहीं है
    • विभिन्न समस्याओं को रैंक की अलग-अलग आवश्यकताएं हैं
  2. समय और मेमोरी जटिलता:
    • FLARE vanilla attention से 200 गुना से अधिक तेज है
    • मेमोरी उपयोग vanilla attention से थोड़ा अधिक है लेकिन Physics Attention से बहुत कम है

आवृत्ति विश्लेषण

O(M³+M²N) समय जटिलता के साथ विशेषता अपघटन एल्गोरिदम के माध्यम से सीखे गए संचार मैट्रिक्स का विश्लेषण:

  • प्रारंभिक ब्लॉक में eigenvalues तेजी से क्षय होते हैं, प्रभावी संपीड़न का संकेत देते हैं
  • गहरे ब्लॉक अधिक अव्यक्त क्षमता का उपयोग करते हैं
  • विभिन्न हेड में विभिन्न आवृत्ति प्रोफाइल होते हैं, स्वतंत्र हेड प्रक्षेपण डिजाइन को सत्यापित करते हैं

संबंधित कार्य

तंत्रिका PDE प्रॉक्सी

  • तंत्रिका ऑपरेटर: FNO, DeepONet आदि अनंत-आयामी फ़ंक्शन स्पेस के बीच मैपिंग सीखते हैं
  • ग्राफ नेटवर्क: जाल पर स्थानीय पड़ोस इंटरैक्शन का उपयोग करते हैं
  • Transformer आर्किटेक्चर: वैश्विक संदर्भ एकत्रीकरण की अनुमति देते हैं लेकिन द्विघात जटिलता से सीमित हैं

कुशल अटेंशन तंत्र

  • Linformer: कुंजी-मूल्य अनुक्रम को प्रक्षेपित करने के लिए रैखिक मैपिंग सीखता है
  • Reformer: स्थानीय संवेदनशील हैशिंग का उपयोग करता है
  • Nyströmformer: स्व-अटेंशन को अनुमानित करने के लिए Nyström विधि का उपयोग करता है
  • LoRA: कुशल सूक्ष्म-ट्यूनिंग के लिए मुख्य रूप से कम-रैंक अनुकूलन

निष्कर्ष और चर्चा

मुख्य निष्कर्ष

  1. FLARE कम-रैंक अटेंशन तंत्र के माध्यम से स्व-अटेंशन की द्विघात जटिलता बाधा को सफलतापूर्वक दरकिनार करता है
  2. कई PDE बेंचमार्क में SOTA सटीकता प्राप्त करता है, साथ ही कम पैरामीटर और कम कम्प्यूटेशनल जटिलता
  3. पहली बार अटेंशन-आधारित तंत्रिका प्रॉक्सी मॉडल को लाखों बिंदु ज्यामिति पर प्रशिक्षण देना संभव बनाता है

सीमाएं

  1. गहरे अवशिष्ट MLP निर्भरता: अनुक्रमिक बाधा और बढ़ी हुई विलंबता का परिचय दे सकता है
  2. निश्चित अव्यक्त टोकन सीमा: M की पसंद को विशिष्ट समस्या के लिए ट्यून करने की आवश्यकता है
  3. कुछ उच्च-रैंक समस्याओं के लिए प्रयोज्यता: जैसे Darcy समस्या में vanilla transformer अभी भी लाभ रखता है

भविष्य की दिशाएं

  1. प्रशिक्षण अवधि के दौरान अव्यक्त टोकन संख्या को क्रमिक रूप से बढ़ाना
  2. प्रसार मॉडलिंग के लिए समय-सशर्त अव्यक्त टोकन डिजाइन करना
  3. स्व-प्रतिगामी मॉडलिंग के लिए केवल-डिकोडर वेरिएंट विकसित करना
  4. गहरे अवशिष्ट MLP की अनुक्रमिक बाधा समस्या को हल करना

गहन मूल्यांकन

शक्तियां

  1. मजबूत तकनीकी नवाचार:
    • अटेंशन राउटिंग समस्या को कम-रैंक मैट्रिक्स अपघटन में चतुराई से रूपांतरित करता है
    • स्वतंत्र हेड प्रक्षेपण डिजाइन विशेष राउटिंग पैटर्न की अनुमति देता है
    • मौजूदा GPU कर्नेल के साथ पूरी तरह संगत
  2. व्यापक प्रयोग:
    • 6 विभिन्न PDE बेंचमार्क को कवर करता है
    • विस्तृत विलोपन प्रयोग और आवृत्ति विश्लेषण
    • पहली बार लाखों बिंदु पैमाने पर प्रयोग
  3. गहन सैद्धांतिक विश्लेषण:
    • O(M³+M²N) समय जटिलता के साथ विशेषता अपघटन एल्गोरिदम प्रदान करता है
    • गणितीय दृष्टिकोण से कम-रैंक संचार की प्रभावशीलता की व्याख्या करता है
    • आवृत्ति विश्लेषण के माध्यम से डिजाइन मान्यताओं को सत्यापित करता है
  4. उच्च व्यावहारिक मूल्य:
    • नया योज्य विनिर्माण डेटासेट जारी करता है
    • कोड ओपन-सोर्स, पुनरुत्पादन में आसान
    • मौजूदा Transformer आर्किटेक्चर में सीधे एकीकृत किया जा सकता है

कमियां

  1. विधि प्रयोज्यता सीमाएं:
    • उच्च-रैंक समस्याओं (जैसे Darcy) पर सीमित प्रभाव
    • M की पसंद को समस्या-विशिष्ट ट्यूनिंग की आवश्यकता है
    • गहरा MLP नई कम्प्यूटेशनल बाधा बन सकता है
  2. प्रयोगात्मक सेटअप सीमाएं:
    • अधिक नई विधियों के साथ तुलना की कमी
    • कुछ बेंचमार्क परीक्षण आकार अपेक्षाकृत छोटे हैं
    • विभिन्न प्रकार की PDE समस्याओं पर सार्वभौमिकता को अधिक सत्यापन की आवश्यकता है
  3. अपर्याप्त सैद्धांतिक विश्लेषण:
    • अभिसरण विश्लेषण की कमी
    • इष्टतम M चयन के लिए सीमित सैद्धांतिक मार्गदर्शन
    • सभी PDE समस्याओं में कम-रैंक मान्यता की तर्कसंगतता को आगे तर्क की आवश्यकता है

प्रभाव

  1. शैक्षणिक योगदान: कुशल अटेंशन तंत्र के लिए नया डिजाइन प्रतिमान प्रदान करता है, विशेष रूप से वैज्ञानिक कंप्यूटिंग क्षेत्र में
  2. व्यावहारिक मूल्य: Transformer को बड़े पैमाने पर ज्यामिति समस्याओं को संभालने में सक्षम बनाता है, AI4Science के विकास को बढ़ावा देता है
  3. पुनरुत्पादनीयता: कोड ओपन-सोर्स, प्रयोगात्मक सेटअप विस्तृत, बाद के अनुसंधान में आसान

लागू परिदृश्य

  • बड़े पैमाने पर अनुरचित जाल पर PDE समाधान
  • बिंदु बादल प्रसंस्करण और ज्यामितीय गहन शिक्षा
  • वैश्विक संचार की आवश्यकता वाली लेकिन कम्प्यूटेशनल संसाधन सीमित अनुक्रम मॉडलिंग कार्य
  • वैज्ञानिक कंप्यूटिंग में प्रॉक्सी मॉडलिंग अनुप्रयोग

संदर्भ

पेपर Transformer, तंत्रिका ऑपरेटर, कुशल अटेंशन तंत्र आदि संबंधित क्षेत्रों के महत्वपूर्ण कार्यों का हवाला देता है, जो इस अनुसंधान के लिए ठोस सैद्धांतिक आधार और तुलना बेंचमार्क प्रदान करता है।


समग्र मूल्यांकन: यह एक उच्च-गुणवत्ता वाला अनुसंधान पेपर है जो Transformer स्केलेबिलिटी समस्या को हल करने के लिए एक नवाचारी समाधान प्रस्तुत करता है। FLARE विधि न केवल सैद्धांतिक रूप से कम-रैंक अपघटन की सुंदर व्याख्या रखती है, बल्कि व्यावहारिक रूप से उत्कृष्ट प्रदर्शन प्रदर्शित करती है। पेपर का प्रयोगात्मक डिजाइन व्यापक है, सैद्धांतिक विश्लेषण गहन है, और बड़े पैमाने पर ज्यामितीय गहन शिक्षा और वैज्ञानिक कंप्यूटिंग को आगे बढ़ाने में महत्वपूर्ण महत्व है।