В батче последовательности дополнены до общей длины. Найдите ошибку в реализации attention: почему она меняет масштаб выхода для коротких последовательностей и как исправить её так, чтобы padded позиции не влияли на нормировку весов? import torch q = torch.randn 2, 4, 8 k = torch.randn 2, 4, 8 v = torch.randn 2, 4, 8 key valid = torch.tensor 1, 1, 1, 1 , 1, 1, 0, 0 , dtype=torch.bool scores = q @ k.transpose 2, 1 weights = torch.softmax scores, dim= 1 weights = weights key valid :, None, : out = weights @ v

Question

Accepted Answer

Ошибка в том, что маска применяется после softmax : padded ключи уже получили часть вероятностной массы. После умножения их веса становятся нулевыми, но веса допустимых ключей не перенормируются, поэтому их сумма становится меньше 1 и выход out уменьшается по масштабу. Маскировать нужно логиты до softmax , присваивая запрещённым ключам −∞ : тогда их вероятность будет ровно нулевой, а вся масса распределится между валидными ключами. Это особенно важно, когда в одном батче длины последовательностей сильно различаются. masked scores = scores.masked fill key valid :, None, : , float " inf" weights = torch.softmax masked scores, dim= 1 out = weights @ v Нужно также гарантировать, что для каждого query есть хотя бы один валидный ключ: softmax от строки, целиком заполненной −∞ , даст неопределённый результат.