Consistency LLM: LLM को parallel decoder में बदलकर inference speed 3.5x बढ़ाना

Consistency Large Language Models (CLLM)

• LLM को पारंपरिक रूप से sequential decoder माना जाता रहा है, जो एक बार में एक token को क्रमिक रूप से decode करते हैं
• इस ब्लॉग में दिखाया गया है कि pre-trained LLM को आसानी से efficient parallel decoder में बदला जा सकता है
• Consistency Large Language Models (CLLM) parallel decoder का एक नया परिवार है, जो हर inference step में n-token sequence को कुशलतापूर्वक decode करके inference latency को कम कर सकता है
• CLLM उस cognitive process की नकल करता है जिसमें इंसान शब्दों को एक-एक करके बोलने से पहले दिमाग में पूरा वाक्य बनाता है, और इसे केवल pre-trained LLM को fine-tuning करके प्रभावी रूप से सिखाया जा सकता है
• CLLM को इस तरह train किया जाता है कि वह random initialization वाले n-token sequence को संभवतः कम से कम steps में autoregressive (AR) decoding के समान परिणाम पर map करे, और यह parallel decoding के जरिए किया जाता है
• प्रयोगों के अनुसार, CLLM ने मौजूदा AR decoder की तुलना में 2.4x से 3.4x तक generation speed improvement दिखाया, और Medusa2 या Eagle जैसी fast inference techniques के बराबर या उनसे बेहतर प्रदर्शन किया
• CLLM यह performance improvement बिना किसी अतिरिक्त memory cost के हासिल कर सकता है

Jacobi decoding की पृष्ठभूमि और सीमाएं

• LLM, AR decoding के जरिए token को एक-एक करके generate करते हैं, इसलिए लंबे response में latency अधिक होती है
• Jacobi decoding, Jacobi और Gauss-Seidel nonlinear equation solvers से आया है, और यह सिद्ध किया गया है कि यह greedy sampling के साथ AR generation के समान है
• Jacobi decoding, sequential generation process को Jacobi iteration आधारित n variables वाले n nonlinear equations के system के रूप में पुनर्गठित करता है, जिससे parallel processing संभव होती है
• हर iteration step में 1 या उससे अधिक सही token का prediction किया जा सकता है, जिससे AR decoding को संभावित रूप से तेज किया जा सकता है
• लेकिन व्यवहार में, AR के लिए train किए गए LLM तब सही token लगभग नहीं बना पाते जब previous tokens में error हो, इसलिए अधिकांश Jacobi iterations में n-token sequence पर केवल एक correction मिलता है, जिससे लंबी Jacobi trajectory बनती है
• Lookahead decoding या speculative decoding, Jacobi decoding की इस inefficiency को कम करने की कोशिश करते हैं, लेकिन inference के समय अतिरिक्त memory cost लाते हैं, जबकि CLLM ऐसा नहीं करता

CLLM की training method

• CLLM training मुख्य रूप से दो हिस्सों में होती है: Jacobi trajectory preparation और consistency तथा AR loss optimization
• Jacobi trajectory preparation चरण में, पूरे response sequence के l tokens generate होने तक उसे n-n के chunks में काटकर क्रमिक रूप से Jacobi decoding की जाती है, और हर trajectory में बने sequence को एक data item माना जाता है
• Training के दौरान consistency loss और AR loss को साथ में optimize किया जाता है; consistency loss यह सुनिश्चित करता है कि कई tokens को एक साथ predict किया जाए, जबकि AR loss यह सुनिश्चित करता है कि CLLM target LLM से बहुत दूर न जाए और generation quality बनी रहे
• Global consistency (GC) loss, Jacobi trajectory के किसी भी बिंदु और fixed point के बीच की दूरी को कम करता है, ताकि CLLM trajectory के किसी भी बिंदु से fixed point predict करने के लिए प्रोत्साहित हो
• Local consistency (LC) loss, Jacobi trajectory की पड़ोसी states को एक जैसा output देने के लिए प्रेरित करता है
• AR loss, target LLM के generation result पर आधारित पारंपरिक AR loss को शामिल करता है, जिसका उद्देश्य CLLM को target LLM के distribution से बाहर जाने से रोकना है

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

• प्रयोगों में Spider (text-to-SQL), Human-Eval (Python code completion), GSM8k (math) जैसे specialized domain tasks और MT-bench जैसे व्यापक open-ended dialogue challenges शामिल थे
• CLLM ने target model की तुलना में सबसे बड़ा speed improvement दिखाया, और inference के समय बिना अतिरिक्त cost के Medusa2 के समान या उससे बेहतर speedup हासिल किया
• MT-bench में, CLLM ने Medusa2 के साथ संयोजन के लगभग समान speed improvement हासिल किया, लेकिन अधिक adaptability और memory efficiency दी
• CLLM की fine-tuning cost मध्यम स्तर की है, और dataset बड़ा होने पर भी Jacobi trajectory generation के लिए dataset के लगभग 10% का उपयोग करके लगभग 2.5x speed improvement पाया जा सका
• CLLM fast forwarding phenomenon के जरिए एक ही Jacobi iteration में कई लगातार tokens को सही ढंग से predict कर सकता है
• CLLM ने stationary token capability दिखाई, जिसमें previous tokens में error होने पर भी वह सही token को पहले से predict करके उसे बदले बिना बनाए रख सकता है
• CLLM training के जरिए collocation जैसे महत्वपूर्ण language concepts सीखता है, और इसके माध्यम से Jacobi trajectory के किसी भी बिंदु पर structure infer करके iteration steps को कम करने के लिए कई शब्दों का एक साथ prediction कर सकता है

GN⁺ की राय

• CLLM ऐसा लगता है कि मौजूदा LLM के AR decoding में मौजूद लंबी latency की समस्या को Jacobi decoding का उपयोग करके प्रभावी रूप से हल करता है। खास तौर पर, अतिरिक्त memory cost के बिना parallelized decoding के जरिए speed improvement हासिल करना प्रभावशाली है

• CLLM की training method, consistency loss के जरिए मौजूदा LLM को fine-tuning करने का अपेक्षाकृत सरल तरीका लगती है, लेकिन इसके माध्यम से language की एक महत्वपूर्ण विशेषता collocation को सीखकर parallel decoding performance को काफी बेहतर बनाना इसे खास बनाता है

• हालांकि, CLLM greedy sampling को मानकर चलता है, इसलिए यह अधिक विविध decoding strategies में भी अच्छी तरह काम करेगा या नहीं, इस पर आगे और शोध की जरूरत दिखती है। साथ ही, अभी के प्रयोग केवल English तक सीमित हैं, इसलिए विभिन्न भाषाओं में इसकी generalization भी सत्यापित की जानी चाहिए

• CLLM, LLM के response speed को बढ़ाने का एक practical approach लगता है। Web search या chatbot जैसे real-time responsiveness वाले tasks में इसका अच्छा उपयोग हो सकता है

• व्यक्तिगत रूप से, CLLM की consistency learning approach से यह उम्मीद बनती है कि क्या इसे GPT जैसे LLM के अलावा image generation model या speech synthesis model जैसे अन्य generative models पर भी लागू किया जा सकता है। उम्मीद है कि CLLM का यह विचार आगे चलकर विभिन्न generative models की efficiency बढ़ाने में योगदान देगा

[GN⁺]

Hacker News राय

• फ्री ड्रॉइंग क्लास में समय सीमा रखकर बार-बार ड्रॉ करने का अभ्यास करने से ड्रॉइंग स्किल काफ़ी बेहतर होती है। जल्दी-जल्दी ड्रॉ करने का अभ्यास उल्टा अनुपात और आउटलाइन को सही पकड़ने की क्षमता बढ़ाता है.
• Jacobi decoding, greedy autoregressive decoding के समान है, लेकिन व्यवहार में iteration से बचने और ज़रूरत से ज़्यादा सामान्य जवाबों से बचने के लिए sampling temperature को 0 से ऊपर सेट करना पड़ता है.
• जल्द ही हमें एहसास होगा कि model training की ज़रूरत नहीं है। सिर्फ़ अच्छी indexing और sampling चाहिए। LLM मूल रूप से एक dataset के DB के बराबर है, जिसमें शानदार NLP interface होता है.
• CLLM पेपर में inference performance का सुधार काफ़ी उल्लेखनीय है। Fine-tuning की लागत वाजिब है (pre-training लागत का लगभग 0.01%) और performance improvement भी काफ़ी consistent है.
• Jacobi trajectories high-temperature paths को बाहर कर देती हैं, इसलिए data retrieval के लिए यह सकारात्मक हो सकता है, लेकिन creativity को अधिकतम करने के लिए नकारात्मक हो सकता है.
• अच्छा होगा अगर "AI experts से पूछें" जैसा कोई feature हो, जो यह समझाए कि LLM एक ही prompt पर deterministic तरीके से जवाब क्यों नहीं देता.
• यह जानना दिलचस्प होगा कि Groq(https://groq.com/) की speedup इसी method की वजह से है या वह कोई अलग तरीका इस्तेमाल करता है.