ലാർജ് ലാംഗ്വേജ് മോഡൽ (LLM) ഇൻഫറൻസിന്റെ പ്രീഫിൽ ഘട്ടത്തിൽ (prefill stage), അറ്റൻഷൻ മാസ്സിന്റെ (attention mass) ഭൂരിഭാഗവും കൈവശം വെക്കുന്ന ടോക്കണുകളെ കണ്ടെത്താൻ പുതിയൊരു “Attention Sink Detector” സഹായിക്കുന്നു. കീ-വാല്യൂ (KV) കാഷെയിൽ നിന്ന് ഈ സിങ്ക് ടോക്കണുകളെ നീക്കം ചെയ്യുന്നത് മോഡലിന്റെ പ്രകടനം തകരാൻ കാരണമാകുമെന്നും, സുസ്ഥിരമായ ജനറേഷന് അവ അത്യന്താപേക്ഷിതമാണെന്നും ഈ ടൂൾ കാണിച്ചുതരുന്നു.

Why the prefill stage matters

മിക്ക LLM ഗവേഷണങ്ങളും ടോക്കൺ തിരിച്ചുള്ള ജനറേഷൻ ലൂപ്പിലാണ് ശ്രദ്ധ കേന്ദ്രീകരിക്കുന്നത്, എന്നാൽ ആദ്യ ടോക്കണിന് മുമ്പ് നടക്കുന്ന പ്രവൃത്തികൾ—പ്രീഫിൽ (prefill)—അടുത്തതായി സംഭവിക്കുന്ന എല്ലാ കാര്യങ്ങളുടെയും സ്വഭാവം നിശ്ചയിക്കുന്നു. പ്രീഫിൽ സമയത്ത് മോഡൽ മുഴുവൻ പ്രോംപ്റ്റും പ്രോസസ്സ് ചെയ്യുകയും, KV കാഷെ നിർമ്മിക്കുകയും, ഓരോ സ്ഥാനത്തും അറ്റൻഷൻ വിതരണം ചെയ്യുകയും ചെയ്യുന്നു. അറ്റൻഷൻ വെയ്റ്റുകൾ ഉത്പാദിപ്പിക്കുന്ന സോഫ്റ്റ്‌മാക്സ് (softmax) തുക 1 ആകേണ്ടതുകൊണ്ട്, "ഒന്നും പ്രസക്തമല്ല" എന്ന് മോഡലിന് പറയാൻ കഴിയില്ല. ബാക്കിയുള്ള പ്രോബബിലിറ്റി മാസ്സ് എളുപ്പത്തിൽ കൈകാര്യം ചെയ്യാവുന്ന ഒരു ടോക്കണിലേക്ക്, സാധാരണയായി സീക്വൻസിലെ ആദ്യത്തെ ടോക്കണിലേക്ക്, മോഡൽ മാറ്റുന്നു. ആ ടോക്കൺ ഒരു attention sink ആയി മാറുന്നു.

What an attention sink is, in plain terms

ഒരു ട്രാൻസ്ഫോർമറിൽ (Transformer), ഓരോ ഹെഡും ഓരോ ടോക്കൺ ജോഡിക്കും ഒരു വെയ്റ്റ് കണക്കാക്കുന്നു. വിതരണം ഒരു വശത്തേക്ക് മാത്രം ചരിഞ്ഞാൽ, ഒരു ടോക്കണിന് അമിതമായ അറ്റൻഷൻ വെയ്റ്റ് ലഭിച്ചേക്കാം. ഈ പ്രതിഭാസം ഒരു പിശകല്ല; അത് സോഫ്റ്റ്‌മാക്സ് നിയന്ത്രണത്തിന്റെ (softmax constraint) ഫലമാണ്. ആദ്യത്തെ ടോക്കൺ (പലപ്പോഴും വാചകത്തിന്റെ തുടക്കത്തിലുള്ള മാർക്കർ) മറ്റ് സ്ഥലങ്ങളിലേക്ക് നൽകാൻ കഴിയാത്ത ബാക്കി പ്രോബബിലിറ്റിക്കുള്ള ഒരു “റിലീസ് വാൽവ്” (release valve) ആയി പ്രവർത്തിക്കുന്നു.

Stakes for KV-cache management

പ്രോസസ്സ് ചെയ്ത ഓരോ ടോക്കണിന്റെയും കീ (key), വാല്യൂ (value) വെക്റ്ററുകൾ KV കാഷെകളിൽ സംഭരിക്കുന്നു, ഇത് ജനറേഷൻ സമയത്ത് വേഗത്തിൽ വിവരങ്ങൾ കണ്ടെത്താൻ സഹായിക്കുന്നു. ദീർഘമായ കോൺടെക്സ്റ്റ് സാഹചര്യങ്ങളിൽ, GPU മെമ്മറി പരിധിക്കുള്ളിൽ നിൽക്കാൻ പ്രായോഗികമായി കാഷെ കുറയ്ക്കാറുണ്ട് (pruning). "അപ്രധാനപ്പെട്ടവ" എന്ന് തോന്നുന്ന ടോക്കണുകളെ അന്ധമായി നീക്കം ചെയ്യുന്നത് മോഡൽ ആശ്രയിക്കുന്ന സിങ്കുകളെയും (sinks) ഇല്ലാതാക്കുമെന്നും ഇത് പെർഫോമൻസ് തകരാൻ കാരണമാകുമെന്നും പുതിയ ഡിറ്റക്ടർ കാണിച്ചുതരുന്നു. കാഷെയിൽ സിങ്ക് ടോക്കണുകൾ നിലനിർത്തുന്നത് മോഡലിന്റെ ആന്തരിക സന്തുലിതാവസ്ഥ നിലനിർത്താനും ജനറേഷൻ സുസ്ഥിരമാക്കാനും സഹായിക്കുന്നു.

How the detector works

  • Eager attention: ഫുൾ അറ്റൻഷൻ മാട്രിക്സിനെ ചുരുക്കുന്ന FlashAttention പോലുള്ള ഫാസ്റ്റ് കർണലുകളെ ഒഴിവാക്കി, പകരം ഓരോ ഹെഡിനും വേണ്ടിയുള്ള റോ (raw) അറ്റൻഷൻ സ്കോറുകൾ ഇത് രേഖപ്പെടുത്തുന്നു.
  • Layer-wide aggregation: ഓരോ ടോക്കണിനും ഒരൊറ്റ അറ്റൻഷൻ-മാസ് പ്രൊഫൈൽ ലഭിക്കുന്നതിനായി എല്ലാ ലെയറുകളിലും ഹെഡുകളിലും സ്കോറുകൾ ശരാശരി കാണുന്നു.
  • Log-median absolute deviation (MAD): അറ്റൻഷൻ വെയ്റ്റുകൾ ഒരു വശത്തേക്ക് കൂടുതൽ ചരിഞ്ഞതുകൊണ്ട്, ലളിതമായ മീൻ-സ്റ്റാൻഡേർഡ്-ഡീവിയേഷൻ (mean-standard-deviation) പരിശോധന സാധാരണ വ്യതിയാനങ്ങളെ ഔട്ട്‌ലയറുകളായി തെറ്റായി തരംതിരിക്കാൻ സാധ്യതയുണ്ട്. വെയ്റ്റുകളെ ലോഗ്-ട്രാൻസ്ഫോം (Log-transform) ചെയ്യുന്നത് വിതരണത്തെ നോർമലൈസ് ചെയ്യുന്നു; തുടർന്ന് MAD പ്രയോഗിക്കുന്നത് സാധാരണ പരിധിയിൽ കൂടുതൽ അറ്റൻഷൻ മാസ് ഉള്ള ടോക്കണുകളെ അടയാളപ്പെടുത്തുന്നു.

Two classes of sinks uncovered

  1. True sinks – പ്രോംപ്റ്റ് ഉള്ളടക്കം എന്തുതന്നെയായാലും, വാചകത്തിന്റെ തുടക്കത്തിലുള്ള (BOS) ടോക്കൺ സ്ഥിരമായി വലിയൊരു ഭാഗം അറ്റൻഷൻ ഉൾക്കൊള്ളുന്നു.
  2. Structural sinks – ChatML-ൽ ഉപയോഗിക്കുന്ന മാർക്കറുകൾ (<im_start>, <im_end>) പോലുള്ള സിസ്റ്റം ലെവൽ പ്രോംപ്റ്റുകളിൽ ഉൾപ്പെടുന്ന ടോക്കണുകൾ അറ്റൻഷൻ ആകർഷിക്കുന്ന ചെറിയ ക്ലസ്റ്ററുകൾ രൂപീകരിക്കുന്നു. അവ ലോജിക്കൽ അതിർവരമ്പുകളായി പ്രവർത്തിക്കുകയും, ഉപയോക്താവിന്റെ സന്ദേശങ്ങളെ സിസ്റ്റം നിർദ്ദേശങ്ങളിൽ നിന്ന് വേർതിരിക്കാൻ മോഡലിനെ സഹായിക്കുകയും ചെയ്യുന്നു.

Counter-point: Do we really need the raw matrix?

റോ ഡാറ്റ (raw numbers) ഇല്ലാതെ സിങ്ക് പ്രതിഭാസം മറഞ്ഞിരിക്കുമെന്നും, അപൂർണ്ണമായ ഡാറ്റയെ അടിസ്ഥാനമാക്കിയുള്ള ഏതൊരു കാഷെ-പ്രുണിംഗ് (cache-pruning) നയവും മോഡലിനെ അസ്ഥിരപ്പെടുത്താൻ സാധ്യതയുണ്ടെന്നും രചയിതാവ് ചൂണ്ടിക്കാട്ടുന്നു.

What to watch next

ജനറേഷൻ ഗുണനിലവാരത്തിനായുള്ള എൻട്രോപ്പി ട്രാക്കറുമായി (entropy tracker) ആരംഭിച്ച നാല് ഭാഗങ്ങളുള്ള പരമ്പരയിലെ രണ്ടാമത്തെ ഭാഗമാണ് ഈ ഡിറ്റക്ടർ. വരാനിരിക്കുന്ന ഭാഗം സിങ്ക് ഡിറ്റക്ഷനെ entropy-guided speculative decoding-മായി സംയോജിപ്പിക്കും. അവസാന പഠനം ഒരു അനുഭവസിദ്ധമായ (empirical) പഠനമായിരിക്കും.

Takeaway: അറ്റൻഷൻ സിങ്കുകൾ വെറുമൊരു പ്രത്യേകതയല്ല; പ്രീഫിൽ സമയത്ത് ഒരു ട്രാൻസ്ഫോർമറിന്റെ അറ്റൻഷൻ വിതരണം കൃത്യമായി നിലനിർത്തുന്ന ഘടനാപരമായ തൂണുകളാണവ. അവ അവഗണിക്കുന്ന ഏതൊരു KV-cache pruning തന്ത്രവും മോഡലിന്റെ സ്ഥിരതയെ അപകടത്തിലാക്കും. ഈ ടോക്കണുകളെ കണ്ടെത്തുകയും സംരക്ഷിക്കുകയും ചെയ്യുന്നത് കുറഞ്ഞ ചിലവിൽ ചെയ്യാവുന്ന ഒരു സുരക്ഷാ ക്രമീകരണമാണ്, ഇത് ദീർഘമായ കോൺടെക്സ്റ്റ് ഇൻഫറൻസിനെ വിശ്വസനീയവും കാര്യക്ഷമവുമാക്കുന്നു.