2025-11-20T08:25:14.880374

Titans: Learning to Memorize at Test Time

Behrouz, Zhong, Mirrokni
Over more than a decade there has been an extensive research effort on how to effectively utilize recurrent models and attention. While recurrent models aim to compress the data into a fixed-size memory (called hidden state), attention allows attending to the entire context window, capturing the direct dependencies of all tokens. This more accurate modeling of dependencies, however, comes with a quadratic cost, limiting the model to a fixed-length context. We present a new neural long-term memory module that learns to memorize historical context and helps attention to attend to the current context while utilizing long past information. We show that this neural memory has the advantage of fast parallelizable training while maintaining a fast inference. From a memory perspective, we argue that attention due to its limited context but accurate dependency modeling performs as a short-term memory, while neural memory due to its ability to memorize the data, acts as a long-term, more persistent, memory. Based on these two modules, we introduce a new family of architectures, called Titans, and present three variants to address how one can effectively incorporate memory into this architecture. Our experimental results on language modeling, common-sense reasoning, genomics, and time series tasks show that Titans are more effective than Transformers and recent modern linear recurrent models. They further can effectively scale to larger than 2M context window size with higher accuracy in needle-in-haystack tasks compared to baselines.
academic

Titans: परीक्षण समय पर स्मरण करना सीखना

मूल जानकारी

  • पेपर ID: 2501.00663
  • शीर्षक: Titans: Learning to Memorize at Test Time
  • लेखक: Ali Behrouz, Peilin Zhong, Vahab Mirrokni (Google Research)
  • वर्गीकरण: cs.LG cs.AI cs.CL
  • प्रकाशन तिथि: 31 दिसंबर 2024
  • पेपर लिंक: https://arxiv.org/abs/2501.00663

सारांश

यह पेपर एक नए तंत्रिका दीर्घकालीन स्मृति मॉड्यूल का प्रस्ताव देता है, जो ऐतिहासिक संदर्भ को स्मरण करना सीख सकता है और ध्यान तंत्र को दीर्घकालीन अतीत की जानकारी का उपयोग करते हुए वर्तमान संदर्भ पर ध्यान केंद्रित करने में मदद करता है। लेखक स्मृति के दृष्टिकोण से तर्क देते हैं कि ध्यान तंत्र सीमित संदर्भ लेकिन सटीक निर्भरता मॉडलिंग के कारण अल्पकालीन स्मृति के रूप में कार्य करता है, जबकि तंत्रिका स्मृति अपनी डेटा स्मरण क्षमता के कारण दीर्घकालीन, अधिक स्थायी स्मृति के रूप में कार्य करती है। इन दोनों मॉड्यूल के आधार पर, लेखक Titans आर्किटेक्चर का एक नया परिवार प्रस्तुत करते हैं और स्मृति को आर्किटेक्चर में प्रभावी ढंग से एकीकृत करने के लिए तीन प्रकार प्रस्तावित करते हैं। प्रायोगिक परिणाम दर्शाते हैं कि Titans भाषा मॉडलिंग, सामान्य ज्ञान तर्क, जीनोमिक्स और समय श्रृंखला कार्यों पर Transformers और आधुनिक रैखिक पुनरावर्ती मॉडल से अधिक प्रभावी है, और 2M से अधिक संदर्भ विंडो आकार तक प्रभावी ढंग से स्केल कर सकता है।

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

मुख्य समस्या

मौजूदा अनुक्रम मॉडलिंग आर्किटेक्चर दक्षता और प्रदर्शन के बीच व्यापार-बंद का सामना करते हैं:

  1. Transformers: हालांकि निर्भरता संबंधों को सटीक रूप से मॉडल कर सकते हैं, लेकिन कम्प्यूटेशनल जटिलता O(n²) है, जो संदर्भ लंबाई को सीमित करती है
  2. रैखिक Transformers/RNNs: हालांकि दक्ष हैं, लेकिन जानकारी को निश्चित आकार की स्थिति में संपीड़ित करते हैं, जिससे लंबे अनुक्रम के प्रदर्शन में गिरावट आती है
  3. स्मृति प्रणाली की कमी: मौजूदा आर्किटेक्चर में मानव मस्तिष्क जैसी बहु-स्तरीय स्मृति प्रणाली (अल्पकालीन स्मृति, दीर्घकालीन स्मृति, मेटा-स्मृति आदि) का अभाव है

अनुसंधान प्रेरणा

लेखक मानव स्मृति प्रणाली से प्रेरित हैं, यह मानते हुए कि प्रभावी सीखने के लिए आवश्यक है:

  • विभिन्न लेकिन परस्पर जुड़े मॉड्यूल, जिनमें से प्रत्येक सीखने की प्रक्रिया में महत्वपूर्ण घटक के लिए जिम्मेदार है
  • डेटा से सक्रिय रूप से सीखने और अतीत के इतिहास के अमूर्तन को स्मरण करने की क्षमता
  • परीक्षण समय पर निरंतर सीखने और अनुकूलन करने का तंत्र

मुख्य योगदान

  1. तंत्रिका दीर्घकालीन स्मृति मॉड्यूल: एक गहन तंत्रिका नेटवर्क को मेटा-मॉडल के रूप में प्रस्तावित किया गया है, जो परीक्षण समय पर अपने पैरामीटर में डेटा को कैसे स्मरण/संग्रहीत करना है यह सीखता है
  2. स्मृति प्रबंधन तंत्र: "आश्चर्य" के आधार पर स्मृति अपडेट तंत्र और अनुकूली विस्मृति तंत्र डिजाइन किया गया है
  3. Titans आर्किटेक्चर परिवार: स्मृति को गहन शिक्षण आर्किटेक्चर में एकीकृत करने के तीन तरीके प्रस्तावित किए गए हैं: Memory as Context (MAC), Memory as Gate (MAG), Memory as Layer (MAL)
  4. समानांतर प्रशिक्षण एल्गोरिदम: तेजी से समानांतर प्रशिक्षण एल्गोरिदम प्रदान किया गया है, जो गहन स्मृति मॉड्यूल प्रशिक्षण को कुशल बनाता है
  5. व्यापक प्रायोगिक सत्यापन: भाषा मॉडलिंग, सामान्य ज्ञान तर्क, जीनोमिक्स और समय श्रृंखला पूर्वानुमान सहित कई कार्यों पर Titans की प्रभावशीलता सत्यापित की गई है

विधि विवरण

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

यह पेपर अनुक्रम मॉडलिंग कार्य का अध्ययन करता है, जहां इनपुट अनुक्रम x∈RN×dinx \in \mathbb{R}^{N \times d_{in}} है, और लक्ष्य एक ऐसा मॉडल सीखना है जो लंबे अनुक्रमों को प्रभावी ढंग से संभाल सके, जिसे आवश्यकता है:

  • परीक्षण समय पर निरंतर सीखने और स्मरण करने की क्षमता
  • अल्पकालीन और दीर्घकालीन स्मृति के उपयोग को संतुलित करना
  • रैखिक जटिलता लेकिन उच्च अभिव्यक्ति क्षमता बनाए रखना

तंत्रिका दीर्घकालीन स्मृति मॉड्यूल

मुख्य डिजाइन विचार

मानव दीर्घकालीन स्मृति से प्रेरित, अप्रत्याशित (आश्चर्यजनक) घटनाएं अधिक आसानी से याद रहती हैं। लेखक इनपुट के सापेक्ष तंत्रिका नेटवर्क के ग्रेडिएंट का उपयोग करके "आश्चर्य" को मापते हैं।

स्मृति अपडेट तंत्र

मूल अपडेट नियम:

M_t = M_{t-1} - θ_t ∇ℓ(M_{t-1}; x_t)

सुधारा गया अपडेट नियम (गति का परिचय):

M_t = M_{t-1} + S_t
S_t = η_t S_{t-1} - θ_t ∇ℓ(M_{t-1}; x_t)

जहां:

  • S_t: आश्चर्य गति, जिसमें अतीत का आश्चर्य और तात्कालिक आश्चर्य शामिल है
  • η_t: डेटा-निर्भर आश्चर्य क्षय पैरामीटर
  • θ_t: तात्कालिक आश्चर्य को एकीकृत करने की डिग्री को नियंत्रित करने वाला पैरामीटर

विस्मृति तंत्र

लंबे अनुक्रमों को संभालने के लिए, अनुकूली विस्मृति तंत्र का परिचय दिया गया है:

M_t = (1 - α_t)M_{t-1} + S_t

जहां α_t ∈ [0,1] एक गेटिंग तंत्र है, जो विस्मृति की डिग्री को नियंत्रित करता है।

उद्देश्य फ़ंक्शन

संबद्ध स्मृति हानि फ़ंक्शन का उपयोग:

ℓ(M_{t-1}; x_t) = ||M_{t-1}(k_t) - v_t||²₂

जहां k_t = x_t W_K, v_t = x_t W_V

Titans आर्किटेक्चर प्रकार

1. Memory as Context (MAC)

  • स्मृति को वर्तमान जानकारी के संदर्भ के रूप में उपयोग करना
  • अनुक्रम को खंडों में संसाधित करना, प्रत्येक खंड दीर्घकालीन स्मृति को प्रासंगिक इतिहास जानकारी प्राप्त करने के लिए क्वेरी करता है
  • ध्यान तंत्र यह तय करता है कि क्या दीर्घकालीन स्मृति जानकारी की आवश्यकता है

2. Memory as Gate (MAG)

  • एक शाखा स्लाइडिंग विंडो ध्यान का उपयोग अल्पकालीन स्मृति के रूप में करती है
  • दूसरी शाखा तंत्रिका स्मृति मॉड्यूल का उपयोग दीर्घकालीन स्मृति के रूप में करती है
  • गेटिंग तंत्र के माध्यम से दोनों शाखाओं के आउटपुट को संयोजित करना

3. Memory as Layer (MAL)

  • तंत्रिका स्मृति को गहन नेटवर्क की एक परत के रूप में उपयोग करना
  • क्रमिक प्रसंस्करण: पहले स्मृति परत के माध्यम से, फिर ध्यान परत के माध्यम से
  • मौजूदा हाइब्रिड मॉडल के समान स्तरीय डिजाइन

समानांतर प्रशिक्षण

प्रशिक्षण प्रक्रिया को मैट्रिक्स गुणन और योग संचालन का उपयोग करके पुनः तैयार करके उच्च-दक्षता समानांतर प्रशिक्षण को लागू किया गया है:

  • अनुक्रम को आकार b के खंडों में विभाजित करना
  • समानांतर संबद्ध स्कैन का उपयोग करके गति शर्तों की गणना करना
  • टेंसर-आधारित छोटे बैच ग्रेडिएंट डिसेंट के माध्यम से तेजी से प्रशिक्षण

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

डेटासेट

  • भाषा मॉडलिंग: FineWeb-Edu डेटासेट, 15B/30B tokens
  • सामान्य ज्ञान तर्क: PIQA, HellaSwag, WinoGrande, ARC-easy/challenge, SIQA, BoolQ
  • लंबे संदर्भ कार्य: RULER benchmark (S-NIAH), BABILong benchmark
  • समय श्रृंखला: ETT, ECL, Traffic, Weather डेटासेट
  • जीनोमिक्स: GenomicsBenchmarks डेटासेट

मॉडल आकार

  • 170M, 340M, 400M, 760M पैरामीटर के मॉडल
  • प्रशिक्षण लंबाई: 4K tokens
  • संदर्भ विंडो: 2M+ tokens तक विस्तारणीय

तुलनात्मक विधियां

  • Transformers: Transformer++
  • रैखिक पुनरावर्ती मॉडल: RetNet, GLA, Mamba, Mamba2, DeltaNet, TTT, Gated DeltaNet
  • हाइब्रिड मॉडल: Samba, Gated DeltaNet-H2
  • बड़े मॉडल: GPT-4, Llama3, RecurrentGemma, Mistral

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

भाषा मॉडलिंग प्रदर्शन

340M पैरामीटर मॉडल पर:

  • Titans (LMM): भ्रम 26.18 (Wiki), 29.97 (LMB)
  • सर्वश्रेष्ठ आधारभूत TTT: भ्रम 27.44 (Wiki), 34.19 (LMB)
  • हाइब्रिड मॉडल में Titans (MAG) सर्वश्रेष्ठ प्रदर्शन: भ्रम 25.07 (Wiki), 28.72 (LMB)

लंबे संदर्भ कार्य

S-NIAH कार्य पर (16K अनुक्रम लंबाई):

  • Titans (MAC): S-NIAH-PK 98.4%, S-NIAH-N 97.4%, S-NIAH-W 95.2%
  • Mamba2: S-NIAH-PK 5.4%, S-NIAH-N 0.0%, S-NIAH-W 0.0%
  • TTT: S-NIAH-PK 88.4%, S-NIAH-N 4.4%, S-NIAH-W 0.0%

BABILong बेंचमार्क परीक्षण

  • Titans कुछ-शॉट सेटिंग में सभी आधारभूत को पार करता है, जिसमें GPT-4 भी शामिल है
  • फाइन-ट्यूनिंग सेटिंग में, छोटे आकार के Titans ने GPT-4 से 70 गुना अधिक पैरामीटर वाले को पार किया

समय श्रृंखला पूर्वानुमान

तंत्रिका स्मृति मॉडल सभी डेटासेट पर आधारभूत विधियों से बेहतर है, जिसमें Mamba, Transformer और रैखिक मॉडल-आधारित विधियां शामिल हैं।

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

घटक योगदान (महत्व के क्रम में):

  1. वजन क्षय (विस्मृति तंत्र)
  2. गति तंत्र
  3. कनवल्शनल परत
  4. स्थायी स्मृति
  5. गहन स्मृति बनाम रैखिक स्मृति

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

रैखिक पुनरावर्ती मॉडल

  • पहली पीढ़ी: RetNet, LRU, RWKV, S4/S5 - डेटा-स्वतंत्र संक्रमण मैट्रिक्स का उपयोग
  • दूसरी पीढ़ी: Griffin, Mamba श्रृंखला - गेटिंग तंत्र का परिचय
  • तीसरी पीढ़ी: DeltaNet, TTT, Longhorn - मेटा-लर्निंग/ऑनलाइन लर्निंग आधारित अपडेट नियम

Transformer प्रकार

  • दक्षता अनुकूलन: विरल ध्यान, रैखिक ध्यान, I/O-जागरूक कार्यान्वयन
  • खंडित Transformer: RMT आदि खंडों के बीच जानकारी पारित करने के लिए सरल वेक्टर स्मृति का उपयोग करते हैं

परीक्षण समय प्रशिक्षण

  • प्रारंभिक स्थानीय सीखने के एल्गोरिदम से प्रेरित
  • MNM और TTT-layer के साथ सबसे अधिक संबंधित, लेकिन Titans में विस्मृति तंत्र और गति अपडेट हैं

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

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

  1. स्मृति प्रणाली का महत्व: बहु-स्तरीय स्मृति प्रणाली (अल्पकालीन + दीर्घकालीन + स्थायी स्मृति) अनुक्रम मॉडलिंग के लिए महत्वपूर्ण है
  2. परीक्षण समय सीखने की प्रभावशीलता: परीक्षण समय पर निरंतर सीखना और स्मरण लंबे अनुक्रम प्रसंस्करण क्षमता को काफी बढ़ा सकता है
  3. आर्किटेक्चर डिजाइन का प्रभाव: MAC और MAG आर्किटेक्चर पारंपरिक MAL परत डिजाइन से बेहतर हैं
  4. विस्तारणीयता सत्यापन: Titans 2M+ संदर्भ विंडो तक प्रभावी ढंग से स्केल कर सकता है

सीमाएं

  1. कम्प्यूटेशनल ओवरहेड: गहन स्मृति मॉड्यूल सरल मैट्रिक्स स्थिति की तुलना में अधिक कम्प्यूटेशनल संसाधन की आवश्यकता है
  2. स्मृति गहराई व्यापार-बंद: गहरे स्मृति मॉड्यूल बेहतर प्रदर्शन करते हैं लेकिन प्रशिक्षण धीमा है
  3. पैरामीटर संवेदनशीलता: आश्चर्य-संबंधित पैरामीटर को सावधानीपूर्वक समायोजित करने की आवश्यकता है
  4. सैद्धांतिक विश्लेषण अपर्याप्त: स्मृति क्षमता और विस्मृति रणनीति के सैद्धांतिक विश्लेषण का अभाव

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

  1. स्मृति आर्किटेक्चर अनुकूलन: अधिक कुशल तंत्रिका स्मृति आर्किटेक्चर डिजाइन की खोज
  2. सैद्धांतिक विश्लेषण: स्मृति क्षमता, विस्मृति रणनीति का सैद्धांतिक विश्लेषण प्रदान करना
  3. बड़े पैमाने पर सत्यापन: बड़े आकार के मॉडल पर विधि की प्रभावशीलता सत्यापित करना
  4. अनुप्रयोग विस्तार: अधिक क्षेत्रों में अनुप्रयोग क्षमता की खोज

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

लाभ

  1. अवधारणा नवाचार मजबूत: मानव स्मृति प्रणाली के दृष्टिकोण से अनुक्रम मॉडलिंग को पुनः सोचना, नए बहु-स्तरीय स्मृति आर्किटेक्चर का प्रस्ताव
  2. तकनीकी योगदान व्यापक: न केवल तंत्रिका स्मृति मॉड्यूल प्रस्तावित किया गया है, बल्कि आर्किटेक्चर में एकीकृत करने के तीन तरीके और उच्च-दक्षता समानांतर एल्गोरिदम डिजाइन किए गए हैं
  3. प्रायोगिक सत्यापन पर्याप्त: कई क्षेत्रों (NLP, समय श्रृंखला, जीनोमिक्स) में व्यापक प्रयोग, परिणाम आश्वस्त करने वाले हैं
  4. सैद्धांतिक आधार ठोस: स्मृति अपडेट को ग्रेडिएंट डिसेंट, गति, वजन क्षय से जोड़ना, सैद्धांतिक व्याख्या प्रदान करना

कमियां

  1. कम्प्यूटेशनल जटिलता विश्लेषण अपर्याप्त: हालांकि रैखिक जटिलता का दावा किया जाता है, लेकिन गहन स्मृति मॉड्यूल की वास्तविक कम्प्यूटेशनल लागत विश्लेषण विस्तृत नहीं है
  2. हाइपरपैरामीटर संवेदनशीलता: कई डेटा-निर्भर पैरामीटर (α_t, θ_t, η_t) की सेटिंग जटिल हो सकती है
  3. मानव स्मृति के साथ सादृश्य सीमित: हालांकि मानव स्मृति से प्रेरित है, लेकिन "आश्चर्य" की परिभाषा अपेक्षाकृत सरल है
  4. बड़े पैमाने के मॉडल सत्यापन अपर्याप्त: सबसे बड़ा मॉडल केवल 760M पैरामीटर है, अरब-स्तरीय पैरामीटर मॉडल सत्यापन की कमी

प्रभाव

  1. शैक्षणिक मूल्य: अनुक्रम मॉडलिंग के लिए नई स्मृति दृष्टिकोण प्रदान करता है, अधिक संबंधित अनुसंधान को प्रेरित कर सकता है
  2. व्यावहारिक मूल्य: लंबे अनुक्रम प्रसंस्करण कार्यों पर उत्कृष्ट प्रदर्शन, व्यावहारिक अनुप्रयोग क्षमता है
  3. पुनरुत्पादनीयता: लेखक कोड ओपन-सोर्स करने का वचन देते हैं, विधि के प्रचार और सत्यापन में सहायता करता है

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

  1. लंबे दस्तावेज प्रसंस्करण: लंबे दस्तावेज प्रसंस्करण की आवश्यकता वाले NLP कार्यों के लिए उपयुक्त
  2. समय श्रृंखला विश्लेषण: विशेष रूप से दीर्घकालीन ऐतिहासिक जानकारी की आवश्यकता वाले पूर्वानुमान कार्यों के लिए उपयुक्त
  3. ऑनलाइन सीखने के परिदृश्य: परीक्षण समय पर निरंतर अनुकूलन की आवश्यकता वाले अनुप्रयोगों के लिए उपयुक्त
  4. स्मृति-गहन कार्य: प्रश्नोत्तर प्रणाली, संवाद प्रणाली आदि जैसे बड़ी मात्रा में जानकारी स्मरण करने की आवश्यकता वाले कार्यों के लिए

संदर्भ

पेपर में 138 संबंधित संदर्भ उद्धृत किए गए हैं, जिसमें Transformer, पुनरावर्ती तंत्रिका नेटवर्क, ध्यान तंत्र, स्मृति नेटवर्क, परीक्षण समय प्रशिक्षण आदि कई संबंधित क्षेत्रों के महत्वपूर्ण कार्य शामिल हैं, जो इस अनुसंधान के लिए ठोस सैद्धांतिक आधार प्रदान करते हैं।