प्रति हेड तीन मैट्रिक्स, या एक बड़ा मैट्रिक्स

प्रति-हेड दृष्टिकोण

प्रत्येक हेड को अपना क्वेरी प्रोजेक्शन, की प्रोजेक्शन, और वैल्यू प्रोजेक्शन चाहिए। हेड h के लिए:

Q_h = X · W_Q^h    जहाँ W_Q^h का शेप [d_model, head_dim] है
K_h = X · W_K^h    जहाँ W_K^h का आकार [d_model, head_dim] है
V_h = X · W_V^h    जहाँ W_V^h का आकार [d_model, head_dim] है

X में इनपुट आकार [batch, seq_len, d_model] होता है। प्रोजेक्शन के बाद, Q_h, K_h, V_h में प्रत्येक का आकार [batch, seq_len, head_dim] होता है।

फ्यूज्ड व्यू

प्रति-हेड मैट्रिक्स मेमोरी में एक-दूसरे के साथ स्थित होते हैं। एक एकल फ्यूज्ड मैट्रिक्स W_Q का आकार [d_model, d_model] सभी हेड्स को एक साथ उत्पन्न करता है:

Q_fused = X · W_Q              # [batch, seq_len, d_model]
Q_per_head = reshape(Q_fused)  # [batch, n_head, seq_len, head_dim]

फ्यूज्ड मैटमल 12 के बजाय एक BLAS कॉल भेजता है। CUDA टेंसर कोर इस आकार के मैटमल पर चरम थ्रूपुट प्राप्त करते हैं; प्रति-हेड मैटमल हार्डवेयर का अपर्याप्त उपयोग करेंगे।

पैरामीटर की संख्या

तीन संयुक्त मैट्रिक्स W_Q, W_K, W_V, प्रत्येक d_model × d_model। प्लस आउटपुट प्रोजेक्शन W_O, भी d_model × d_model। ANDREA-120M के लिए:

प्रति लेयर के अटेंशन के params = 4 × 768² = 2,359,296 ≈ 2.36M
12 लेयर्स में params      = 12 × 2.36M ≈ 28.3M

ANDREA-120M के कुल पैरामीटर्स का लगभग एक-चौथाई हिस्सा अटेंशन प्रोजेक्शन्स में रहता है। बाकी तीन-चौथाई हिस्सा MLP सबलेयर और एम्बेडिंग्स में रहता है।

प्रोजेक्शन्स का नामकरण

बारह वेक्टर्स एक बन जाते हैं

प्रत्येक हेड के गणना करने के बाद

हर हेड का आउटपुट टेंसर का आकार [batch, seq_len, head_dim] होता है। बारह हेड बारह ऐसे टेंसर उत्पन्न करते हैं। फीचर आयाम के साथ संयोजन (Concatenation) उन्हें वापस एक साथ जोड़ देता है:

concat_output = concat(head_1, head_2, ..., head_12)
shape         = [batch, seq_len, n_head × head_dim]
= [batch, seq_len, 768]    # for ANDREA-120M

Concat विघटन को उलट देता है। कुल फीचर आयाम d_model पर वापस आ जाता है। आयामों में कोई सूचना हानि नहीं; अंतर इस बात में है कि प्रत्येक चंक में क्या शामिल है: हेड 1 का चंक हेड 1 के सीखे हुए ध्यान पैटर्न को प्रतिबिंबित करता है।

आउटपुट प्रोजेक्शन W_O

केवल संयोजन हेड्स को अलग-थलग छोड़ देता है: हेड 4 का आउटपुट हेड 7 के आउटपुट के बगल में बैठा होता है, न तो एक-दूसरे के बारे में जागरूक। आउटपुट प्रोजेक्शन W_O आकार [d_model, d_model] के आकार का उन्हें मिलाता है:

attention_output = concat_output · W_O
आकार            = [batch, seq_len, d_model]

W_O के बाद, हर आउटपुट आयाम में बारहों हेड्स का एक सीखा रैखिक संयोजन होता है। जानकारी इस एकल मैट्रिक्स गुणा के माध्यम से हेड्स के बीच स्वतंत्र रूप से बहती है।

हेड्स क्यों विशेषज्ञता प्राप्त करते हैं

आर्किटेक्चर में कुछ भी हेड 4 को क्रिया काल सीखने या हेड 9 को जोड़ी गई विराम चिह्नों को सीखने के लिए मजबूर नहीं करता। विशेषज्ञता ग्रेडिएंट दबाव से उभरती है: प्रशिक्षण के दौरान, जो हेड्स अतिशयोक्तिपूर्ण योगदान देते हैं वे उन हेड्स से छोटे ग्रेडिएंट्स प्राप्त करते हैं जो अद्वितीय योगदान देते हैं। हजारों चरणों में, प्रत्येक हेड कुल हानि को सबसे प्रभावी ढंग से कम करने वाली एक निचे में बस जाता है।

प्रयोगात्मक रूप से, प्रशिक्षित ट्रांसफॉर्मर्स में ऐसे हेड्स दिखाई देते हैं जो संभालते हैं: स्थिति पैटर्न (हेड पिछले टोकन को देखता है), सिंटैक्टिक पैटर्न (हेड मिलते हुए बंद ब्रैकेट को देखता है), सेमांटिक पैटर्न (हेड सबसे हाल के नामित इकाई को देखता है)। कोई लेबल इस विशेषज्ञता को प्रशिक्षित नहीं करते। प्रशिक्षण सिग्नल अकेले, W_O के माध्यम से प्रचारित, हेड्स को वर्गीकृत करता है।

बारह हेड्स क्यों, एक चौड़ा हेड नहीं

CUDA हेड्स को कैसे स्टोर करता है

एक ही टेंसर, पुनः आकारित

ANDREA का ट्रेनिंग इंजन microgpt_cuda.cu बारह हेड्स के लिए बारह अलग-अलग बफर्स आवंटित नहीं करता। यह एक फ्यूज्ड टेंसर आवंटित करता है & हेड आयाम को स्ट्राइड पैटर्न के रूप में मानता है:

// Q = X · W_Q के बाद (एक मैटमल, हेड्स में फ्यूज्ड)
// Q का आकार [batch, seq_len, d_model] है
// आकार बदलें    [batch, seq_len, n_head, head_dim]
// स्थानांतरित करें  [batch, n_head, seq_len, head_dim]
// अब प्रत्येक head के आंतरिक दो आयामों में सन्निकट (contiguous) हैं

स्थानांतरण n_head को seq_len से पहले ले जाता है। क्यों? क्योंकि अगला ऑपरेशन (Q_h · K_h^T) को प्रत्येक head का seq_len × head_dim स्लाइस स्मृति में सन्निकट (contiguous) चाहिए। CUDA matmuls सन्निकट टेंसर्स पर तेज़ चलते हैं।

एक Kernel, कई Heads

एक ही अटेंशन CUDA कर्नेल सभी हेड्स पर समानांतर रूप से चलता है। प्रत्येक थ्रेड ब्लॉक एक (बैच, हेड) जोड़ी को हैंडल करता है; ब्लॉक के अंदर के थ्रेड्स seq_len × head_dim टाइल पर सहयोग करते हैं। कर्नेल को कभी पता नहीं चलता कि वह कई हेड्स को प्रोसेस कर रहा है; लॉन्च ग्रिड समानांतरता को हैंडल करता है।

कॉन्फ़िगरेशन हार्डवेयर को प्रतिबिंबित करता है

ANDREA-120M का n_head=12, head_dim=64 का चयन RTX 4090 टेंसर कोर के साथ संरेखित होता है, जो 16 के गुणकों में matmul टाइल्स को प्राथमिकता देते हैं। head_dim=64 = 4 × 16 टाइल आकार से ठीक मेल खाता है। head_dim=32 (ANDREA-12M) भी मेल खाता है लेकिन टाइल का अपूर्ण उपयोग करता है। head_dim=72 मेल नहीं खाएगा & फॉलबैक कर्नेल्स को मजबूर करेगा।

अंतिम चित्र

पांच चरण। तीन मैटमल इनपुट को छूते हैं (Q, K, V प्रोजेक्शन्स)। एक मैटमल संयुक्त हेड्स को छूता है (W_O)। एक ध्यान कर्नेल सभी हेड्स को समानांतर में संभालता है। ANDREA-120M प्रत्येक लेयर प्रति एक बार सभी पांच चरण चलाता है, बारह लेयर गहरी, हर फॉरवर्ड पास में।

पाइपलाइन पढ़ना