Attention Sink Weights

Problem

Inject an attention sink score into every row and softmax over the sink plus the original scores, so each output row has one extra leading weight for the sink.

The signature is attention_sink_softmax(scores, sink_score). Use a stable max over the sink and the row; masked future columns are handled separately by the caller.

Hint

The sink behaves like one extra key with its own fixed score that every query can always attend to.

Test cases

case 1

in: [[[0]],0]

expected: [[0.5,0.5]]

case 2

in: [[[1,2]],2]

expected: [[0.4223187983,0.1553624035,0.4223187983]]

case 3

in: [[[0,0],[1000,1000]],1000]

expected: [[1,0,0],[0.3333333333,0.3333333333,0.3333333333]]

case 4

in: [[[-5]],0]

expected: [[0.9933071491,0.0066928509]]

Starter codePython

import math
def attention_sink_softmax(scores, sink_score):
    # Your code here
    pass

4 test cases · run your code against them in the editor