Wiki заметки
func.func @attention_head(
%x : tensor<seq x d_model>,
%wq : tensor<d_model x d_k>,
%wk : tensor<d_model x d_k>,
%wv : tensor<d_model x d_v>) {
...
}
Возьмём следующие размеры:
- seq = 4 (длина последовательности 4 токена)
- d_model = 8 (скрытое пространство 8)
- d_k = 4 (размерность ключей/запросов)
- d_v = 4 (размерность значений)
func.func @mini_attention(
%x : tensor<4x8>, // 4 tokens, 8 vectors to define embedding
%wq : tensor<8x4>, // 8 vectors to define embedding,
%wk : tensor<8x4>, // 8 vectors to define embedding
%wv : tensor<8x4> // 8 vectors to define embedding
) -> tensor<4x4> {
...
// 2 tokena могут корелировать но ты не можешь узнть на сколько из их эмбэддинга
// чтобы узнать насколько они корелируют нужно дополнительное пространство которое обучено на то чтобы детектировать кореляции
// чтобы получить это пространство мы умножаем текущий эмбэдинг на wQ и wK
// Q of each = embedding x wQ
%Q = linalg.matmul
ins(%x, %wq : tensor<4x8>, tensor<8x4>)
outs(%q_init : tensor<4x4>) -> tensor<4x4>
// K of each = embedding x wK
%K = linalg.matmul
ins(%x, %wk : tensor<4x8>, tensor<8x4>)
outs(%k_init : tensor<4x4>) -> tensor<4x4>
// тут также для того чтобы корелировать-умножать свапаем размерности
//%K_t = tensor.transpose %K, [1, 0] : tensor<4x4> -> tensor<4x4>
// как сильно что влияет
%scores = linalg.matmul
ins(%Q, %K : tensor<4x4>, tensor<4x4>)
outs(%scores_init : tensor<4x4>) -> tensor<4x4>
// нормализация
%attn = linalg.softmax
ins(%scores : tensor<4x4>)
{axis = 1} -> tensor<4x4>
// дальше нужно сдвинть екущий эмбэддинг в векторном пространстве чз сильно корелирующие тоены/эмбэдинги
// для этого нужно получить пространство которое позволит расчитать как повлияет
%V = linalg.matmul
ins(%x, %wv : tensor<4x8>, tensor<8x4>)
outs(%v_init : tensor<4x4>) -> tensor<4x4>
// Output = attn * V
// сдвигаем эмбэдинг в адресном пространстве
%Y = linalg.matmul
ins(%attn, %V : tensor<4x4>, tensor<4x4>)
outs(%y_init : tensor<4x4>) -> tensor<4x4>
}
def solution(arg):
if arg > 10:
return 42
return 0
I am trying to convert ML Operations like Convolution (which uses 4 nested loops) to High Level Operations (HLO) using xla-legalize-tf, followed by converting HLO to linalg and linalg to affine loops as seen from the above diagram.
random.shuffle(docs) //перемешивает элементы списка `docs` случайным образом прямо на месте**
// Объединяет элементы списка `docs` в одну строку
docs = ["abc", "de", "f"]
"".join(docs) # "abcdef"
// Преобразует строку в **множество уникальных символов**.
set("abcdef")
# {'a', 'b', 'c', 'd', 'e', 'f'}
// matrix x vector multiply
// vector - столбец, каждая строка матрицы умножается на столбец-вектор
def linear(x, w):
return [sum(wi * xi for wi, xi in zip(wo, x)) for wo in w]
/*
Матрица
W
1 2 3
4 5 6
Вектор
x
[10
20
30
]
Результат:
[1⋅10+2⋅20+3⋅30
4⋅10+5⋅20+6⋅30]
[1⋅10+2⋅20+3⋅30
4⋅10+5⋅20+6⋅30]
*/
//Softmax должен выдать:
// числа > 0
// сумма = 1
// Экспонента → нормализация делением → готово.
def softmax(logits):
// find max
max_val = max(val.data for val in logits)
// exp(val - max_val)
// экспонент нормализованных логитов
exps = [(val - max_val).exp() for val in logits]
total = sum(exps)
return [e / total for e in exps]
1. Зачем нужно `(val - max_val)`?
Без этого:
+ большие значения могут вызвать переполнение (exp(1000) → слишком большое число) - numerical stability trick
logits = [3.0, 1.0, -2.0]
max_val = 3.0
val | val - max_val | exp(val - max_val)
3 0 1.0
1 -2 0.1353…
-2 -5 0.0067…
NOTE:
Логит (logit) — это сырое, не нормализованное выходное значение нейронной сети до применения softmax или сигмоиды.
2. зачем exponent:
+ Преобразует логиты в положительные значения
+ Усиливает различия между логитами
логит exp(logit)
1 2.7
2 7.4
3 20
Разница между 1 и 3 была ×3,
после exp стала ×7,5.
Это делает модель более «решительной».
немного больший логит → намного большая вероятность
намного меньший логит → почти 0 вероятность
+ Softmax хорошо работает с градиентами:
функция гладкая
производная легко вычисляется
обучение стабильнее
3. Softmax на NPU - в 3 шага
+ Reduce-max — нахождение максимума
MPU HW делит вектор logits на блоки и на каждом блоке параллельно ищет максимум
выполняется за O(n / параллелизм-на-блоках)
+ Subtract+Exp - вычитание максимума и экспонента
val_i = logits[i] - max_val - параллельно на всех ядрах
exp - на таблице значений - LUT (таблица предвычисленных значений - exponential lookup table)
+ Reduce-sum и деление
sum_exp = Σ exp(val_i) - in parralel - чз дерево суммирования - это бинарная структура данных,
используемая для эффективного параллельного вычисления суммы массива за время, где
— число элементов. В этом дереве каждый узел равен сумме своих двух дочерних узлов,
что позволяет объединять результаты вычислений на каждом уровне параллельно,
начиная с листьев и заканчивая корнем.
деление
p_i = exp(val_i) / sum_exp
+ 5 очень быстрых векторных инструкций
reduce_max
vector_sub
vector_exp
reduce_sum
vector_div