FFT की वापसी: Self-Attention का एक कुशल विकल्प

 (arxiv.org)
3 पॉइंट द्वारा GN⁺ 2025-02-27 | 1 टिप्पणियां | WhatsApp पर शेयर करें
  • लंबे context वाले Transformer में self-attention cost bottleneck बनने की स्थिति में, SPECTRE FFT-आधारित token mixer के जरिए प्रति layer complexity को O(L²) से O(L log L) तक घटाता है
  • हर attention head को fast real FFT, content-adaptive spectral gate, और inverse FFT के संयोजन से बदला जाता है, जबकि मौजूदा Transformer structure बना रहता है
  • autoregressive generation में Prefix-FFT cache हर step पर FFT दोबारा compute करने का बोझ घटाता है, और optional wavelet module local feature loss की भरपाई कर सकता है
  • Llama-3.2-1B backbone पर SDPA, FlashAttention-2, और SPECTRE की तुलना की गई, तथा NVIDIA A100-80GB पर 512~128k token throughput और latency मापी गई
  • SPECTRE ने PG-19 और ImageNet-1k पर baseline के बराबर या उससे बेहतर performance दिखाई, और 6% से कम अतिरिक्त parameters के साथ सामान्य GPU पर लंबे context processing को लक्ष्य बनाया

FFT से Self-attention की quadratic cost घटाने का तरीका

  • लंबे context वाले Transformer की ज़रूरत multi-turn dialogue, किताब-लंबाई summary, और high-resolution vision जैसे दसियों हज़ार tokens वाले tasks में होती है
  • मौजूदा self-attention में O(n²d) cost के कारण context लंबा होने पर inference latency और memory usage बढ़ जाते हैं
  • SPECTRE self-attention layer को frequency-domain token mixer से बदलने वाला एक drop-in replacement तरीका है
    • tokens को orthonormal Fourier basis में project किया जाता है
    • content-adaptive diagonal gate और optional low-rank gate लागू किए जाते हैं
    • inverse transform से फिर token space में लौटाया जाता है
  • आसपास के network architecture को बदले बिना प्रति layer complexity को O(n log n) तक घटाना इसका मुख्य बिंदु है

Token mixer की संरचना और generation support

  • SPECTRE के attention head replacement में fast real FFT, spectral gate, और inverse FFT शामिल हैं
  • spectral gating n/2 + 1 frequency coefficients पर काम करती है, ताकि computation और memory usage घटे लेकिन expressiveness बनी रहे
  • Prefix-FFT cache standard KV-cache जैसा काम करता है और streaming decoding को support देता है
    • autoregressive generation में हर time step पर FFT दोबारा compute करने की मौजूदा spectral mixer की कमजोरी को कम करता है
    • fixed memory budget के भीतर efficient generation संभव बनाता है
  • optional Wavelet Refinement Module शुद्ध spectral approach में खो सकने वाले local details की भरपाई करता है, और इसका computational overhead छोटा है

मौजूदा Transformer में इसे लागू करने का तरीका

  • SPECTRE multi-head attention layer को सीधे replace कर सकता है, इसलिए अलग architecture redesign की ज़रूरत नहीं होती
  • मौजूदा pre-trained model को SPECTRE layer के साथ fine-tune किया जा सकता है
    • update का लक्ष्य नए जोड़े गए parameters होते हैं
    • अतिरिक्त parameters कुल weights के 6% से कम हैं
  • specialized optimization या non-standard architecture की ज़रूरत वाले approaches के विपरीत, यह आसपास की Transformer structure को बनाए रखता है

Llama-3.2-1B आधारित प्रयोग

  • एक ही Llama-3.2-1B backbone पर तीन attention kernels लगाकर तुलना की गई
    • standard softmax-dot-product attention(SDPA)
    • FlashAttention-2
    • SPECTRE mixer
  • मापन environment NVIDIA A100-80GB था, और sequence length L ∈ {512, 1k, 4k, 8k, 32k, 128k} थी
  • metrics थे tokens-per-second throughput और single-batch latency
    • throughput जितना अधिक, उतना बेहतर
    • latency जितनी कम, उतना बेहतर
  • SPECTRE ने backbone accuracy बनाए रखते हुए लगभग O(n log n) के करीब runtime दिखाया
    • 32k tokens तक runtime लगभग flat बना रहा
    • abstract के अनुसार 128k-token context में FlashAttention-2 से अधिकतम 7× तेज
    • main text की contribution list के अनुसार 32k tokens पर FlashAttention-2 से अधिकतम 7× तेज inference

Benchmark results और practical scope

  • SPECTRE ने PG-19 language modeling और ImageNet-1k classification में baseline के बराबर या उससे बेहतर results दिखाए
  • लंबे context processing में self-attention की quadratic cost से बचते हुए भी global context mixing बनाए रखता है
  • sparse pattern, kernel approximation, और low-rank structure-आधारित attention acceleration methods में exactness की कमी, non-standard optimization, या streaming generation support न होने जैसी सीमाएँ हो सकती हैं
  • SPECTRE frequency-domain approach इस्तेमाल करता है, जहाँ FFT circular convolution को diagonalize करके global mixing को element-wise product में बदल देता है
  • अतिरिक्त parameters को 6% से कम रखकर, specialized hardware के बिना commodity GPU पर hundred-kilotoken context processing को लक्ष्य बनाया गया है

संबंधित पढ़ाई

1 टिप्पणियां

 
GN⁺ 2025-02-27
Hacker News की राय

  • मूल रूप से यह convolution theorem का उपयोग करने का तरीका है: मूल space में महंगा convolution, reciprocal space में सिर्फ multiplication बन जाता है, और उल्टा भी सही है
    अगर data में convolution operation है, तो उसे conjugate domain में transform करके multiplication में बदल दें
    दूसरे शब्दों में, data के लिए natural domain में काम करें
    https://en.wikipedia.org/wiki/Convolution_theorem

    • इस तरह कहना बहुत अच्छा है, लेकिन LLM में structured attention space का frequency domain होना मेरे लिए बिल्कुल obvious नहीं था
    • यह बुनियादी mathematical space transform sandwich है: 1) data को दूसरे space में बदलें, 2) उस space में operation करें, फिर 3) मूल space में वापस लाएं
      optimize करना हो तो हर step को optimize करें, और जहां संभव हो सबसे efficient space में ज्यादा काम करें
    • “data के लिए natural domain में काम करें” वाली बात में, मुझे समझ नहीं आता कि multiplication को convolution की तुलना में किसी domain में ज्यादा natural क्यों माना जाए
      क्या यह सिर्फ computation आसान होने से अलग बात नहीं है?
    • क्या reciprocal space हमेशा frequency = 1/time की तरह बस 1/space के रूप में होता है?
    • सही है, लेकिन बचत काफी हद तक theoretical है। O(n²) operation को O(nlog n) में बदलना अच्छा लगता है, जब तक आपको यह एहसास न हो कि average n 3 है
      ऊपर से calculation में complex numbers इस्तेमाल करने पड़ते हैं, और numerically भी कम stable होता है। मेरी जानकारी में FFT सामान्य convolution में फायदा नहीं देता
      self-attention या इस paper के use case में n बहुत बड़ा हो सकता है। paper नहीं पढ़ा है। फिर भी complex numbers वाली समस्या बनी रहती है

  • Google ने 2022 में FNet: Mixing Tokens with Fourier Transforms के साथ यह idea पेश किया था
    बाद में पता चला कि ज्यादातर स्थितियों में TPU की matrix multiplication performance FFT से तेज है
    https://arxiv.org/abs/2105.03824

    • इस paper में भी इसका citation है:
      “कुल मिलाकर FNet, Performer, sparse transformer जैसे approaches दिखाते हैं कि fixed या approximate token mixing से computational burden घटाया जा सकता है, लेकिन हमारी adaptive spectral filtering strategy FFT की efficiency को learnable और input-dependent spectral filters के साथ uniquely combine करती है। यह complex sequence modeling tasks के लिए महत्वपूर्ण scalability और adaptability का मजबूत combination देती है।”
      उसके बाद comparison section भी है
    • specialized hardware बेहतर है, यह comparison थोड़ा अजीब लगता है
      वैसे DSP में FFT की मदद करने वाला dedicated hardware होता है क्या? सच में जिज्ञासा से पूछ रहा हूं। इस्तेमाल तो नहीं किया है, लेकिन धुंधला-सा लगता है कि मदद मिल सकती है
    • GPU ने TPU की तुलना में 10% improvement दिखाया
      “TPU Fourier transform में इतने inefficient हैं कि researchers ने 4096 से कम sequences के लिए FFT algorithm इस्तेमाल नहीं किया, बल्कि precomputed DFT matrix का उपयोग करने वाला quadratic-scaling Fourier transform implementation चुना।”
      “Nvidia Quadro P6000 GPU पर FNet architecture में Fourier transform ने inference time का अधिकतम 30% लिया।”
      इस company ने 2021 में दावा किया था कि अगर Google TPU में उनके optical chips इस्तेमाल करे, तो inference time 40% घट सकता है। अगर FFTNet ज्यादा काम संभाले, तो यह और घट सकता है
      https://scribe.rip/optalysys/attention-fourier-transforms-a-...
    • जैसे-जैसे context window में tokens की संख्या बढ़ाते हैं, FFT scaling और बेहतर होती लगती है। Google models का context size में competitors से आगे होना दिलचस्प है
    • सिर्फ FFT से तेज होने की बात नहीं है; TPU का FFT support हमेशा best-effort level का रहा है। आखिरी बार जब कोशिश की थी, तो गंभीर precision issues थे

  • Fourier transform “token” dimension के साथ apply होता है। लेकिन कई applications में इस dimension का कोई अर्थ नहीं होता। इसलिए transformer permutation-invariant data handle करने के लिए अच्छा विकल्प बनता है
    कम चर्चित finite groups पर Fourier transform का उपयोग करके और experiments देखना चाहूंगा। यह permutation-invariant है, फिर भी standard Fourier transform के साथ कई properties साझा करता है
    और अगर यह LLM की अगली बड़ी wave बनता है, तो vLLM या llama.cpp जैसे inference engines इसे कितनी आसानी से integrate कर पाएंगे, यह भी जानना चाहूंगा
    https://en.wikipedia.org/wiki/Fourier_transform_on_finite_gr...

    • मैं इस field का expert नहीं हूं, लेकिन क्या ज्यादातर models में tokens position-dependent information के साथ transform नहीं होते?
      मेरी समझ में llama input में position के हिसाब से vector पर rotation apply करता है
    • इस case में finite group क्या है?

  • गणित मेरे सिर के ऊपर से निकल जाता है, और सूत्रों के आसपास की व्याख्या भी बस मुश्किल से समझ आती है। क्या कोई आसान भाषा में समझा सकता है कि यह attention mechanism के बराबर कैसे है?
    यहां frequency से क्या मतलब है, और tokens के बीच positional relationships कैसे encode किए जाते हैं?

    • Fourier transform एक invertible operator है। यानी यह functions पर काम करता है, और matrices के मामले में function और operator दोनों को matrix के रूप में व्यक्त किया जा सकता है। यह उन्हें उस जगह में transform करता है जिसे हम frequency space कहते हैं
      signal analysis या images में यह सबसे intuitive होता है: https://homepages.inf.ed.ac.uk/rbf/HIPR2/fourier.htm
      frequency space मूलतः complex numbers से व्यक्त किया गया “complex” space है। frequency का फायदा यह है कि यह problem को global तरीके से देखती है
      यह mechanism attention mechanism के बराबर नहीं है, और इसमें साफ़ trade-off है। हालांकि attention जिन महत्वपूर्ण relationships को पकड़ता है, उनमें से काफी को पकड़ने की संभावना जरूर है
      modReLU के बारे में फिलहाल मेरे पास अच्छी intuition नहीं है, लेकिन frequency को modify करते हुए भी inverse Fourier transform को preserve करने की वजह से यह महत्वपूर्ण लगता है
    • असली mechanism अपने-आप में काफी simple है। input embeddings पर FFT apply करता है, फिर input embeddings से MLP द्वारा मिले weights के साथ element-wise multiplication करता है, उसके बाद constant लेकिन learnable bias जोड़ता है, activation function से गुजारता है, और आखिर में inverse FFT apply करता है
      यहां “frequency” शायद काफी abstract चीज हो सकती है। FFT अक्सर ऐसे तरीकों से भी इस्तेमाल होता है जहां कोई साफ़ frequency interpretation नहीं होती। convolution theorem जैसी सुविधाजनक mathematical properties की वजह से इसका इस्तेमाल कई बार किया जाता है
      अगर यह सच में अच्छी तरह काम करता है तो यह काफी चौंकाने वाला और बहुत elegant है
    • मैं बिल्कुल expert नहीं हूं, लेकिन थोड़ी intuition जोड़ूं तो self-attention आखिरकार एक parameterized token mixer है
      यानी output का हर vector इस बात पर निर्भर करता है कि संबंधित input vector बाकी सभी input vectors के किसी function द्वारा कैसे transform हुआ है
      https://medium.com/optalysys/attention-fourier-transforms-a-...
      conceptually देखा जा सकता है कि यह थोड़े simplified convolution जैसा कैसे है: https://openreview.net/pdf?id=8l5GjEqGiRG
      convolution अक्सर तब इस्तेमाल होता है जब किसी न किसी तरीके से global state को ध्यान में रखना हो

  • इस framework में causal masking डालने के लिए लगता है n अलग-अलग FFT करने पड़ेंगे, और positional embeddings का भी कोई जिक्र नहीं है
    इसलिए comparison वाली self-attention implementation शायद non-causal NoPE है, और अगर ऐसा है तो यह baseline को जानबूझकर कमजोर रखने का मामला हो सकता है, इसलिए इतना impressive नहीं
    अगर results state-of-the-art के करीब होते तो लेखक ने शायद इसका जिक्र किया होता

    • Long Range Arena(LRA) benchmark में वे दिखाते जरूर हैं कि उनका model सभी categories में जीतता है। उम्मीद है उन्होंने वे categories या बेहतर models exclude नहीं किए होंगे जहां वे हारे

  • यह related reference लगता है: https://arxiv.org/abs/2111.13587
    Adaptive Fourier Neural Operators: Efficient Token Mixers for Transformers
    John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, Bryan Catanzaro

  • मैं जानना चाहता हूं कि यहां frequency domain में देखने से मदद क्यों मिलती है, इसकी कोई intuition है क्या
    DC component समझ आता है, लेकिन मुझे उम्मीद नहीं है कि input data इतना periodic होगा कि दूसरी frequencies meaningful हों

  • लगता है कि कुछ साल पहले ही O(n log n) full-context mixing दिखाने वाली Hyena Operator की prior work का जिक्र नहीं है
    https://arxiv.org/abs/2302.10866

  • Big O notation से कुछ हद तक अंदाज़ा तो मिल जाता है, लेकिन computer science या electrical engineering से जुड़ी ज़्यादातर चीज़ों की तरह यह भी मेरे सिर के ऊपर से निकल जाती है
    गणित में सचमुच कमजोर होने के नाते, मुझे उन लोगों से ईर्ष्या होती है जो ऐसी चीज़ों को समझ सकते हैं, या कम से कम सीखकर engineering degree और license तक हासिल कर लेते हैं
    FFT के बारे में मुझे बस इतना पता है कि यह signals को बदलता है, किसी तरह की signal processing में इस्तेमाल होता है, और मैंने सुना था कि पुराने समय में nuclear explosion detection के लिए यह मुख्य तकनीक था

    • Fourier transform के बारे में ठीक-ठाक intuition होना बहुत उपयोगी tool है, भले ही आप हाथ से Fourier transform derive न कर सकें या FFT algorithm खुद न लिख सकें
      मूल idea यह है: लगभग हर उपयोगी signal को अलग-अलग frequency और phase वाली sine waves के योग के रूप में व्यक्त किया जा सकता है। उदाहरण के लिए electric signal या sound wave 1D signal हैं, जिनमें x-axis time होता है। देखने में यह संभालने में कठिन, जटिल टेढ़ी-मेढ़ी line हो सकती है
      Fourier transform का उपयोग करके time-based signal की अलग-अलग frequencies को अलग किया जा सकता है। फिर आप किसी खास frequency को अपनी इच्छा के अनुसार modify कर सकते हैं। जैसे अगर signal में बहुत सारा random नुकीला noise है, तो वह high frequency के रूप में दिखाई देता है। इसे साफ़ करने के लिए Fourier transform करें, किसी तय threshold से ऊपर की frequency का data हटा दें, फिर बचे हुए data पर inverse Fourier transform लगाकर original signal के अधिक smooth version में लौट आएँ। इसे low-pass filter कहते हैं, और यह original signal का moving average लेने जैसा ही है
      दिलचस्प बात यह है कि इसे काफ़ी intuitive तरीके से higher dimensions में बढ़ाया जा सकता है। x-axis और y-axis दोनों spatial हों तो 2D signal एक image होता है। JPEG compression इसी concept पर आधारित है। image को छोटा store करने के लिए high-frequency signal हटाए जाते हैं, और बदले में fine details खो जाती हैं, या बहुत ज़्यादा हटाने पर ring-shaped artifacts बनते हैं। इसमें time नाम का तीसरा dimension जोड़ दें तो video बन जाता है, और इसे आगे भी बढ़ाया जा सकता है
      यह सब visually समझना अच्छा है, इसलिए पूरी गहराई से math जाने बिना भी अच्छी intuition मिल सकती है। visualization और interactive examples से भरा एक अच्छा page: https://www.jezzamon.com/fourier/index.html
      3Blue1Brown का video भी अच्छी तरह समझाता है: https://youtu.be/spUNpyF58BY?si=dz0z-s8NftW3Htun
    • सरल शब्दों में, मान लीजिए आपके पास 1D time-domain signal है, जैसे mic से measure किया गया audio signal। अगर mic fixed है, तो वह किसी खास point पर time के साथ air displacement को measure कर रहा है
      Fourier transform, जिसका discrete version FFT है, उस 1D time-domain signal को frequency बनाम magnitude और phase components में तोड़ देता है
      frequency मूल रूप से pitch है। pure sine wave या pure tone वैसी आवाज़ जैसी होती है जो पहले देर रात TV broadcast बंद होते समय सुनाई देती थी; इस case में ज़्यादातर values 0 होती हैं और उस tone की frequency location पर एक “spike” बनता है। signal amplitude जितना बड़ा होगा, spike का size भी उतना बड़ा होगा। pitch, यानी frequency, ऊपर या नीचे जाती है तो यह spike horizontal axis के साथ इधर-उधर खिसकता है
      phase मूल रूप से signal का time offset है। किसी भी तरह delay किया गया tone अलग phase के रूप में दिखाई देता है। हालांकि यह absolute measurement नहीं, relative measurement है। unit radians यानी angle होती है, इसलिए circle का एक पूरा चक्कर लगाने पर यह फिर “reset” हो जाता है; इसलिए यह पता नहीं चल सकता कि signal 1 second delay हुआ या 2 seconds
      इसलिए एक signal, यानी time बनाम amplitude से, वास्तव में आपको दो तरह की जानकारी मिलती है: frequency बनाम magnitude और phase
      अगर आप imaginary numbers या complex variables समझते हैं, तो ये दो signals असल में FFT output, जो एक complex function है, की magnitude और argument ही हैं

  • Telemetry के दौर में, cloud telemetry पर FFT apply करके periodic anomalies और quasi-stable systems को incident होने के बाद नहीं बल्कि पहले पहचानना, एक बड़ा मौका गंवाने जैसा लगता है
    दुर्भाग्य से, यह मेरे पहचान पाने के level पर तो है, लेकिन implement करने की skill level पर नहीं, और मेरा schedule पहले से ही भरा हुआ है
    “SLA सबसे ज़्यादा service deploy होने के 23–25 minutes बाद violate होता है। हम्म, ऐसा क्यों होगा… ओह नहीं।”

    • “माफ़ कीजिए Dave, आपकी application deploy नहीं की जा सकती”
      मज़ाक छोड़ें तो, जिस जगह यह सच में पैसा बना सकता है वह है traffic cycles का prediction करके server instances को ऊपर-नीचे करना और cost घटाना
      ऐसी चीज़ अगर personal time में की जाए तो company कभी approve नहीं करेगी, लेकिन अगर इसे ready-made product के रूप में package किया जाए तो company तुरंत खरीद लेगी