ਇੱਕ ਨਵਾਂ “Attention Sink Detector” ਉਹਨਾਂ ਟੋਕਨਾਂ ਦੀ ਪਛਾਣ ਕਰਦਾ ਹੈ ਜੋ ਲਾਰਜ-ਲੈਂਗੂਏਜ-ਮਾਡਲ (LLM) ਇਨਫਰੈਂਸ ਦੇ ਪ੍ਰੀਫਿਲ (prefill) ਸਟੇਜ ਦੌਰਾਨ ਜ਼ਿਆਦਾਤਰ ਅਟੈਂਸ਼ਨ ਮਾਸ (attention mass) ਨੂੰ ਇਕੱਠਾ ਕਰਦੇ ਹਨ। ਇਹ ਟੂਲ ਦਿਖਾਉਂਦਾ ਹੈ ਕਿ ਕੀ-ਵੈਲਯੂ (KV) ਕੈਸ਼ ਤੋਂ ਇਹਨਾਂ ਸਿੰਕ ਟੋਕਨਾਂ ਨੂੰ ਹਟਾਉਣ ਨਾਲ ਮਾਡਲ ਦੀ ਕਾਰਗੁਜ਼ਾਰੀ ਖਰਾਬ ਹੋ ਜਾਂਦੀ ਹੈ, ਜੋ ਉਹਨਾਂ ਨੂੰ ਸਥਿਰ ਜਨਰੇਸ਼ਨ ਲਈ ਜ਼ਰੂਰੀ ਐਂਕਰ ਬਣਾਉਂਦਾ ਹੈ।
ਪ੍ਰੀਫਿਲ ਸਟੇਜ ਕਿਉਂ ਮਹੱਤਵਪੂਰਨ ਹੈ
ਜ਼ਿਆਦਾਤਰ LLM ਖੋਜ ਟੋਕਨ-ਦਰ-ਟੋਕਨ ਜਨਰੇਸ਼ਨ ਲੂਪ 'ਤੇ ਕੇਂਦਰਿਤ ਹੁੰਦੀ ਹੈ, ਪਰ ਪਹਿਲੇ ਟੋਕਨ ਤੋਂ ਪਹਿਲਾਂ ਹੋਣ ਵਾਲਾ ਕੰਮ—ਪ੍ਰੀਫਿਲ—ਅਗਲੇ ਸਾਰੇ ਕੰਮਾਂ ਲਈ ਅਧਾਰ ਤਿਆਰ ਕਰਦਾ ਹੈ। ਪ੍ਰੀਫਿਲ ਦੌਰਾਨ ਮਾਡਲ ਪੂਰੇ ਪ੍ਰੋਂਪਟ ਨੂੰ ਪ੍ਰੋਸੈਸ ਕਰਦਾ ਹੈ, KV ਕੈਸ਼ ਬਣਾਉਂਦਾ ਹੈ, ਅਤੇ ਹਰ ਸਥਾਨ 'ਤੇ ਅਟੈਂਸ਼ਨ ਫੈਲਾਉਂਦਾ ਹੈ। ਕਿਉਂਕਿ ਸਾਫਟਮੈਕਸ (softmax) ਜੋ ਅਟੈਂਸ਼ਨ ਵੇਟਸ ਪੈਦਾ ਕਰਦਾ ਹੈ, ਉਸਦਾ ਜੋੜ 1 ਹੋਣਾ ਚਾਹੀਦਾ ਹੈ, ਮਾਡਲ ਇਹ ਨਹੀਂ ਕਹਿ ਸਕਦਾ ਕਿ "ਕੁਝ ਵੀ ਪ੍ਰਸੰਗਿਕ ਨਹੀਂ ਹੈ।" ਇਹ ਬਚੀ ਹੋਈ ਪ੍ਰੋਬੇਬਿਲਟੀ ਮਾਸ (probability mass) ਨੂੰ ਇੱਕ ਅਜਿਹੇ ਟੋਕਨ 'ਤੇ ਸੁੱਟ ਦਿੰਦਾ ਹੈ ਜਿਸ ਨੂੰ ਸੰਬੋਧਿਤ ਕਰਨਾ ਆਸਾਨ ਹੋਵੇ, ਜੋ ਆਮ ਤੌਰ 'ਤੇ ਸੀਕੁਇੰਸ ਵਿੱਚ ਪਹਿਲਾ ਟੋਕਨ ਹੁੰਦਾ ਹੈ। ਉਹ ਟੋਕਨ ਇੱਕ ਅਟੈਂਸ਼ਨ ਸਿੰਕ (attention sink) ਬਣ ਜਾਂਦਾ ਹੈ।
ਸਰਲ ਸ਼ਬਦਾਂ ਵਿੱਚ ਅਟੈਂਸ਼ਨ ਸਿੰਕ ਕੀ ਹੈ
ਇੱਕ Transformer ਵਿੱਚ, ਹਰੇਕ ਹੈੱਡ (head) ਹਰ ਟੋਕਨ ਜੋੜ ਲਈ ਇੱਕ ਵੇਟ (weight) ਦੀ ਗਣਨਾ ਕਰਦਾ ਹੈ। ਜਦੋਂ ਡਿਸਟ੍ਰੀਬਿਊਸ਼ਨ ਬਹੁਤ ਜ਼ਿਆਦਾ ਇੱਕ ਪਾਸੇ ਝੁਕ ਜਾਂਦੀ ਹੈ, ਤਾਂ ਇੱਕ ਸਿੰਗਲ ਟੋਕਨ ਨੂੰ ਅਟੈਂਸ਼ਨ ਵੇਟ ਦਾ ਅਨੁਪਾਤਕ ਹਿੱਸਾ ਮਿਲ ਸਕਦਾ ਹੈ। ਇਹ ਵਰਤਾਰੇ ਕੋਈ ਬੱਗ (bug) ਨਹੀਂ ਹੈ; ਇਹ ਸਾਫਟਮੈਕਸ ਦੀ ਪਾਬੰਦੀ ਕਾਰਨ ਹੁੰਦਾ ਹੈ। ਪਹਿਲਾ ਟੋਕਨ (ਅਕਸਰ ਵਾਕ ਦੀ ਸ਼ੁਰੂਆਤ ਦਾ ਮਾਰਕਰ) ਬਚੀ ਹੋਈ ਪ੍ਰੋਬੇਬਿਲਟੀ ਲਈ ਇੱਕ “ਰਿਲੀਜ਼ ਵਾਲਵ” (release valve) ਵਜੋਂ ਕੰਮ ਕਰਦਾ ਹੈ ਜਿਸ ਨੂੰ ਕਿਤੇ ਹੋਰ ਨਿਰਧਾਰਤ ਨਹੀਂ ਕੀਤਾ ਜਾ ਸਕਦਾ।
KV-ਕੈਸ਼ ਪ੍ਰਬੰਧਨ ਲਈ ਜੋਖਮ
KV ਕੈਸ਼ ਹਰੇਕ ਪ੍ਰੋਸੈਸ ਕੀਤੇ ਗਏ ਟੋਕਨ ਲਈ ਕੀ (key) ਅਤੇ ਵੈਲਯੂ (value) ਵੈਕਟਰਾਂ ਨੂੰ ਸਟੋਰ ਕਰਦੇ ਹਨ, ਜੋ ਜਨਰੇਸ਼ਨ ਦੌਰਾਨ ਤੇਜ਼ ਲੁੱਕ-ਅੱਪ ਕਰਨ ਦੇ ਯੋਗ ਬਣਾਉਂਦੇ ਹਨ। ਲੰਬੇ-ਕੰਟੈਕਸ (long-context) ਵਾਲੇ ਮੈਟਰ ਵਿੱਚ, ਪ੍ਰੈਕਟੀਸ਼ਨਰ GPU ਮੈਮੋਰੀ ਦੀਆਂ ਸੀਮਾਵਾਂ ਦੇ ਅੰਦਰ ਰਹਿਣ ਲਈ ਕੈਸ਼ ਨੂੰ ਪ੍ਰੂਨ (prune) ਕਰਦੇ ਹਨ। ਨਵਾਂ ਡਿਟੈਕਟਰ ਦਿਖਾਉਂਦਾ ਹੈ ਕਿ ਬਿਨਾਂ ਕਿਸੇ ਸੋਚ-ਵਿਚਾਰ ਦੇ ਉਹਨਾਂ ਟੋਕਨਾਂ ਨੂੰ ਡਿਲੀਟ ਕਰਨਾ ਜੋ “ਅਣਮਹੱਤਵਪੂਰਨ” ਲੱਗਦੇ ਹਨ, ਉਹਨਾਂ ਸਿੰਕਸ ਨੂੰ ਵੀ ਡਿਲੀਟ ਕਰ ਦਿੰਦਾ ਹੈ ਜਿਨ੍ਹਾਂ 'ਤੇ ਮਾਡਲ ਨਿਰਭਰ ਕਰਦਾ ਹੈ, ਜਿਸ ਨਾਲ ਕਾਰਗੁਜ਼ਾਰੀ ਡਿੱਗ ਜਾਂਦੀ ਹੈ। ਕੈਸ਼ ਵਿੱਚ ਸਿੰਕ ਟੋਕਨਾਂ ਨੂੰ ਰੱਖਣਾ ਮਾਡਲ ਦੇ ਅੰਦਰੂਨੀ ਸੰਤੁਲਨ ਨੂੰ ਬਣਾਈ ਰੱਖਦਾ ਹੈ ਅਤੇ ਜਨਰੇਸ਼ਨ ਨੂੰ ਸਥਿਰ ਰੱਖਦਾ ਹੈ।
ਡਿਟੈਕਟਰ ਕਿਵੇਂ ਕੰਮ ਕਰਦਾ ਹੈ
- Eager attention: ਇਹ ਇੰਪਲੀਮੈਂਟੇਸ਼ਨ FlashAttention ਵਰਗੇ ਫਾਸਟ ਕਰਨਲਸ (fast kernels) ਨੂੰ ਬਾਈਪਾਸ ਕਰਦੀ ਹੈ, ਜੋ ਪੂਰੇ ਅਟੈਂਸ਼ਨ ਮੈਟ੍ਰਿਕਸ ਨੂੰ ਇਕੱਠਾ ਕਰ ਦਿੰਦੇ ਹਨ, ਅਤੇ ਇਸ ਦੀ ਬਜਾਏ ਹਰੇਕ ਹੈੱਡ ਲਈ ਰੌਅ ਅਟੈਂਸ਼ਨ ਸਕੋਰ ਰਿਕਾਰਡ ਕਰਦੀ ਹੈ।
- Layer-wide aggregation: ਹਰੇਕ ਟੋਕਨ ਲਈ ਇੱਕ ਸਿੰਗਲ ਅਟੈਂਸ਼ਨ-ਮਾਸ ਪ੍ਰੋਫਾਈਲ ਤਿਆਰ ਕਰਨ ਲਈ ਸਾਰੇ ਲੇਅਰਾਂ ਅਤੇ ਹੈੱਡਾਂ ਵਿੱਚ ਸਕੋਰਾਂ ਦਾ ਔਸਤ ਕੱਢਿਆ ਜਾਂਦਾ ਹੈ।
- Log-median absolute deviation (MAD): ਕਿਉਂਕਿ ਅਟੈਂਸ਼ਨ ਵੇਟਸ ਬਹੁਤ ਜ਼ਿਆਦਾ ਰਾਈਟ-ਸਕਿਊਡ (right-skewed) ਹੁੰਦੇ ਹਨ, ਇੱਕ ਸਧਾਰਨ ਮੀਨ-ਸਟੈਂਡਰਡ-ਡਿਵੀਏਸ਼ਨ ਟੈਸਟ ਆਮ ਵੈਰੀਐਂਸ ਨੂੰ ਆਊਟਲਾਈਅਰ ਵਜੋਂ ਗਲਤ ਸ਼੍ਰੇਣੀ ਵਿੱਚ ਪਾ ਸਕਦਾ ਹੈ। ਵੇਟਸ ਨੂੰ ਲੌਗ-ਟਰਾਂਸਫਾਰਮ ਕਰਨਾ ਡਿਸਟ੍ਰੀਬਿਊਸ਼ਨ ਨੂੰ ਨਾਰਮਲਾਈਜ਼ ਕਰਦਾ ਹੈ; ਫਿਰ MAD ਲਾਗੂ ਕਰਨ ਨਾਲ ਉਹਨਾਂ ਟੋਕਨਾਂ ਨੂੰ ਫਲੈਗ ਕੀਤਾ ਜਾਂਦਾ ਹੈ ਜਿਨ੍ਹਾਂ ਦਾ ਅਟੈਂਸ਼ਨ ਮਾਸ ਆਮ ਫੈਲਾਅ ਤੋਂ ਵੱਧ ਹੁੰਦਾ ਹੈ।
ਖੋਜੇ ਗਏ ਸਿੰਕਸ ਦੀਆਂ ਦੋ ਸ਼੍ਰੇਣੀਆਂ
- True sinks – ਵਾਕ ਦੀ ਸ਼ੁਰੂਆਤ (BOS) ਵਾਲਾ ਟੋਕਨ ਪ੍ਰੋਂਪਟ ਦੀ ਸਮੱਗਰੀ ਦੀ ਪਰਵਾਹ ਕੀਤੇ ਬਿਨਾਂ ਲਗਾਤਾਰ ਅਟੈਂਸ਼ਨ ਦਾ ਇੱਕ ਵੱਡਾ ਹਿੱਸਾ ਸੋਖ ਲੈਂਦਾ ਹੈ।
- Structural sinks – ਉਹ ਟੋਕਨ ਜੋ ਸਿਸਟਮ-ਲੇਵਲ ਦੇ ਪ੍ਰੋਂਪਟ ਨਾਲ ਸਬੰਧਤ ਹੁੰਦੇ ਹਨ, ਜਿਵੇਂ ਕਿ ChatML ਵਿੱਚ ਵਰਤੇ ਜਾਣ ਵਾਲੇ ਮਾਰਕਰ (
<im_start>,<im_end>), ਛੋਟੇ ਕਲਸਟਰ ਬਣਾਉਂਦੇ ਹਨ ਜੋ ਅਟੈਂਸ਼ਨ ਨੂੰ ਵੀ ਆਕਰਸ਼ਿਤ ਕਰਦੇ ਹਨ। ਉਹ ਤਰਕਿਕ ਸੀਮਾਵਾਂ ਵਜੋਂ ਕੰਮ ਕਰਦੇ ਹਨ, ਜੋ ਮਾਡਲ ਨੂੰ ਯੂਜ਼ਰ ਮੈਸੇਜ ਨੂੰ ਸਿਸਟਮ ਹਦਾਇਤਾਂ ਤੋਂ ਵੱਖ ਕਰਨ ਵਿੱਚ ਮਦਦ ਕਰਦੇ ਹਨ।
ਵਿਰੋਧੀ ਨੁਕਤਾ: ਕੀ ਸਾਨੂੰ ਸੱਚਮੁੱਚ ਰੌਅ ਮੈਟ੍ਰਿਕਸ (raw matrix) ਦੀ ਲੋੜ ਹੈ?
ਲੇਖਕ ਇਹ ਨੁਕਤਾ ਉਠਾਉਂਦਾ ਹੈ ਕਿ ਰੌਅ ਨੰਬਰਾਂ ਤੋਂ ਬਿਨਾਂ ਸਿੰਕ ਵਰਤਾਰਾ ਲੁਕਿਆ ਰਹਿੰਦਾ ਹੈ, ਅਤੇ ਅਧੂਰੇ ਡੇਟਾ 'ਤੇ ਅਧਾਰਤ ਕੋਈ ਵੀ ਕੈਸ਼-ਪ੍ਰੂਨਿੰਗ ਨੀਤੀ ਮਾਡਲ ਨੂੰ ਅਸਥਿਰ ਕਰਨ ਦਾ ਜੋਖਮ ਲੈਂਦੀ ਹੈ।
ਅੱਗੇ ਕੀ ਦੇਖਣਾ ਹੈ
ਇਹ ਡਿਟੈਕਟਰ ਚਾਰ-ਭਾਗਾਂ ਵਾਲੀ ਲੜੀ ਦਾ ਦੂਜਾ ਹਿੱਸਾ ਹੈ ਜੋ ਜਨਰੇਸ਼ਨ ਕੁਆਲਿਟੀ ਲਈ ਐਂਟ੍ਰੋਪੀ ਟ੍ਰੈਕਰ ਨਾਲ ਸ਼ੁਰੂ ਹੋਈ ਸੀ। ਆਉਣ ਵਾਲਾ ਹਿੱਸਾ ਸਿੰਕ ਡਿਟੈਕਸ਼ਨ ਨੂੰ entropy-guided speculative decoding ਨਾਲ ਜੋੜੇਗਾ। ਅੰਤਿਮ ਅਧਿਐਨ ਇੱਕ ਅਨੁਭਵਿਕ (empirical) ਅਧਿਐਨ ਹੋਵੇਗਾ।
ਮੁੱਖ ਗੱਲ (Takeaway): ਅਟੈਂਸ਼ਨ ਸਿੰਕਸ ਕੋਈ ਅਜੀਬ ਵਿਸ਼ੇਸ਼ਤਾ ਨਹੀਂ ਹਨ; ਇਹ ਉਹ ਸੰਰਚਨਾਤਮਕ ਥੰਮ੍ਹ ਹਨ ਜੋ ਪ੍ਰੀਫਿਲ ਦੌਰਾਨ ਇੱਕ Transformer ਦੇ ਅਟੈਂਸ਼ਨ ਡਿਸਟ੍ਰੀਬਿਊਸ਼ਨ ਨੂੰ ਸਹੀ ਰੱਖਦੇ ਹਨ। ਕੋਈ ਵੀ KV-ਕੈਸ਼ ਪ੍ਰੂਨਿੰਗ ਰਣਨੀਤੀ ਜੋ ਉਹਨਾਂ ਨੂੰ ਨਜ਼ਰਅੰਦਾਜ਼ ਕਰਦੀ ਹੈ, ਉਹ ਮਾਡਲ ਦੀ ਸਥਿਰਤਾ ਨੂੰ ਖਤਰੇ ਵਿੱਚ ਪਾਵੇਗੀ। ਇਹਨਾਂ ਟੋਕਨਾਂ ਦੀ ਪਛਾਣ ਕਰਨਾ ਅਤੇ ਉਹਨਾਂ ਨੂੰ ਬਚਾਉਣਾ ਇੱਕ ਘੱਟ ਲਾਗਤ ਵਾਲਾ ਸੁਰੱਖਿਆ ਉਪਾਅ ਹੈ ਜੋ ਲੰਬੇ-ਕੰਟੈਕਸ ਇਨਫਰੈਂਸ ਨੂੰ ਭਰੋਸੇਯੋਗ ਅਤੇ ਕੁਸ਼ਲ ਬਣਾ ਸਕਦਾ ਹੈ।