QK-Norm Attention Scores

Problem

Apply QK-Norm: RMS-normalize every query and key row, then compute scaled dot-product scores with the usual 1 / sqrt(d) factor. This removes row magnitude from the dot products and keeps logits bounded.

The signature is qk_norm_scores(Q, K, eps=0.0). Return the (seq_q, seq_k) score matrix.

Hint

After RMS normalization the mean square of each row is 1, so self-scores scale with the dimension as sqrt(d) times 1 / sqrt(d).

Test cases

case 1

in: [[[2,0],[0,3]],[[1,0],[0,1]],0]

expected: [[1.4142135624,0],[0,1.4142135624]]

case 2

in: [[[1,0]],[[1,0]],0]

expected: [[1.4142135624]]

case 3

in: [[[3,4]],[[3,4]],0]

expected: [[1.4142135624]]

case 4

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

expected: [[1.1313708499]]

Starter codePython

import math
def qk_norm_scores(Q, K, eps=0.0):
    # Your code here
    pass

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