Error: todos los pesos de atención suman 1

«El softmax normaliza, así que toda la matriz de pesos suma 1.»

scores = np.array([[2., 1., 0.], [0., 2., 1.]])
rowwise = np.exp(scores) / np.exp(scores).sum(axis=-1, keepdims=True)
assert np.allclose(rowwise.sum(axis=-1), 1.0)
assert np.isclose(rowwise.sum(), 2.0)

Cada fila suma 1 y la matriz suma 2: una por consulta. El softmax corre sobre el último eje, así que cada fila es su propia distribución sobre las claves.

Normalizar toda la matriz también da 1 y es algo completamente distinto, por lo que el total general es una mala comprobación.

Para un tensor de puntajes (lote, cabezas, consultas, claves), el softmax toma el último eje, el de las claves.