નવું “Attention Sink Detector” એવા ટોકન્સને ઓળખે છે જે LLM ઇન્ફરન્સના prefill stage દરમિયાન મોટાભાગના attention mass ને એકઠા કરે છે. આ ટૂલ દર્શાવે છે કે KV cache માંથી આ sink tokens ને દૂર કરવાથી મોડેલનું પ્રદર્શન બગડી જાય છે, જે તેમને સ્થિર જનરેશન માટે આવશ્યક એન્કર બનાવે છે.

prefill stage શા માટે મહત્વનું છે

મોટાભાગના LLM સંશોધનો token-by-token જનરેશન લૂપ પર ધ્યાન કેન્દ્રિત કરે છે, પરંતુ પ્રથમ ટોકન પહેલા જે કામ થાય છે—prefill—તે પછીના તમામ કાર્યો માટે પાયો નાખે છે. Prefill દરમિયાન મોડેલ સમગ્ર પ્રોમ્પ્ટને પ્રોસેસ કરે છે, KV cache બનાવે છે, અને દરેક પોઝિશન પર attention ફેલાવે છે. કારણ કે attention weights ઉત્પન્ન કરતા softmax નો સરવાળો 1 થવો જોઈએ, મોડેલ એમ કહી શકતું નથી કે "કંઈ પણ સુસંગત નથી." તે બાકી રહેલા પ્રોબેબિલિટી માસને એવા ટોકન પર નાખે છે જેને એડ્રેસ કરવું સરળ હોય, જે સામાન્ય રીતે સિક્વન્સમાં પ્રથમ ટોકન હોય છે. તે ટોકન attention sink બની જાય છે.

સાદી ભાષામાં attention sink શું છે

Transformer માં, દરેક હેડ દરેક ટોકન જોડી માટે વજન (weight) ની ગણતરી કરે છે. જ્યારે વિતરણ (distribution) ખૂબ જ અસંતુલિત હોય, ત્યારે એક જ ટોકનને attention weight નો અપ્રમાણસર હિસ્સો મળી શકે છે. આ ઘટના કોઈ બગ (bug) નથી; તે softmax નિયમનું પરિણામ છે. પ્રથમ ટોકન (જે ઘણીવાર beginning-of-sentence માર્કર હોય છે) બાકી રહેલી પ્રોબેબિલિટી માટે "release valve" તરીકે કામ કરે છે જેને અન્ય ક્યાંય અસાઇન કરી શકાતી નથી.

KV-cache મેનેજમેન્ટ માટેના જોખમો

KV caches પ્રોસેસ થયેલા દરેક ટોકન માટે key અને value vectors સ્ટોર કરે છે, જે જનરેશન દરમિયાન ઝડપી લુક-અપ્સ સક્ષમ બનાવે છે. લાંબા કોન્ટેક્સ્ટ (long-context) ના કિસ્સામાં, પ્રેક્ટિશનર્સ GPU મેમરી મર્યાદામાં રહેવા માટે કેશનું pruning કરે છે. નવું ડિટેક્ટર દર્શાવે છે કે જે ટોકન્સ "અગત્યના" દેખાય છે તેમને અંધાધૂંધ ડિલીટ કરવાથી તે જ sinks પણ ડિલીટ થઈ જાય છે જેના પર મોડેલ નિર્ભર છે, જેના કારણે પ્રદર્શન તૂટી જાય છે. કેશમાં sink tokens ને પિન કરવાથી મોડેલનું આંતરિક સંતુલન જળવાય છે અને જનરેશન સ્થિર રહે છે.

ડિટેક્ટર કેવી રીતે કામ કરે છે

  • Eager attention: આ અમલીકરણ FlashAttention જેવા ફાસ્ટ કર્નલ્સને બાયપાસ કરે છે, જે સંપૂર્ણ attention મેટ્રિક્સને સંકોચે છે, અને તેના બદલે દરેક હેડ માટે રો (raw) attention સ્કોર્સ રેકોર્ડ કરે છે.
  • Layer-wide aggregation: દરેક ટોકન દીઠ સિંગલ attention-mass પ્રોફાઇલ બનાવવા માટે તમામ લેયર્સ અને હેડ્સમાં સ્કોર્સનું સરેરાશ (average) કાઢવામાં આવે છે.
  • Log-median absolute deviation (MAD): કારણ કે attention weights ખૂબ જ right-skewed હોય છે, સાદો mean-standard-deviation ટેસ્ટ સામાન્ય વેરિઅન્સને આઉટલાયર્સ તરીકે ખોટી રીતે વર્ગીકૃત કરે છે. વજનને log-transform કરવાથી વિતરણ નોર્મલાઇઝ થાય છે; ત્યારબાદ MAD લાગુ કરવાથી એવા ટોકન્સને ફ્લેગ કરવામાં આવે છે જેનું attention mass સામાન્ય વિતરણ કરતા વધારે હોય છે.

શોધાયેલા sinks ના બે પ્રકાર

  1. True sinks – પ્રોમ્પ્ટના કન્ટેન્ટને ધ્યાનમાં લીધા વિના, beginning-of-sentence (BOS) ટોકન સતત એટેન્શનનો મોટો હિસ્સો શોષી લે છે.
  2. Structural sinks – જે ટોકન્સ સિસ્ટમ-લેવલના પ્રોમ્પ્ટ્સના ભાગ છે, જેમ કે ChatML માં વપરાતા માર્કર (<im_start>, <im_end>), તે નાના ક્લસ્ટર્સ બનાવે છે જે એટેન્શનને આકર્ષે છે. તેઓ લોજિકલ બોર્ડર તરીકે કામ કરે છે, જે મોડેલને યુઝર મેસેજને સિસ્ટમ ઇન્સ્ટ્રક્શનથી અલગ કરવામાં મદદ કરે છે.

વિરોધ પક્ષ: શું આપણને ખરેખર raw મેટ્રિક્સની જરૂર છે?

લેખક જણાવે છે કે રો (raw) નંબર્સ વગર sink ઘટના છુપાયેલી રહે છે, અને અધૂરા ડેટા પર આધારિત કોઈપણ કેશ-પ્રુનિંગ પોલિસી મોડેલને અસ્થિર કરવાનું જોખમ ઊભું કરે છે.

આગળ શું જોવું

આ ડિટેક્ટર ચાર ભાગની શ્રેણીનો બીજો ભાગ છે જે જનરેશન ક્વોલિટી માટે એન્ટ્રોપી ટ્રેકરથી શરૂ થઈ હતી. આગામી ભાગમાં sink ડિટેક્શનને entropy-guided speculative decoding સાથે જોડવામાં આવશે. અંતિમ અભ્યાસ એક empirical study હશે.

Takeaway: Attention sinks એ કોઈ અસ્પષ્ટ વિચિત્રતા નથી; તેઓ માળખાગત સ્તંભો છે જે prefill દરમિયાન Transformer ના attention distribution ને વ્યવસ્થિત રાખે છે. કોઈપણ KV-cache pruning વ્યૂહરચના જે તેમને અવગણે છે તે મોડેલની સ્થિરતા જોખમમાં મૂકશે. આ ટોકન્સને શોધવા અને સાચવવાનો એ એક ઓછી કિંમતનો સુરક્ષા ઉપાય છે જે લાંબા કોન્ટેક્સ્ટ ઇન્ફરન્સને વિશ્વસનીય અને કાર્યક્ષમ બનાવી શકે છે.