🔥 Команда дня: einsum или как реализовать multi-head self-attention без единого цикла

Если вы работаете с нейросетями, особенно с трансформерами, то, скорее всего, сталкивались с реализациями self-attention, переполненными циклами. Однако благодаря np.einsum можно выразить всю механику multi-head attention в компактной и векторизованной форме.

Вот пример реализации:

def multi_head_attention(X, W_q, W_k, W_v, W_o):  
    d_k = W_k.shape[-1]  
    Q = np.einsum('si,hij->hsj', X, W_q)  # (n_heads, seq_len, d_k)
    K = np.einsum('si,hik->hsk', X, W_k)  
    V = np.einsum('si,hiv->hsv', X, W_v)  
    scores = Q @ K.transpose(0, 2, 1) / np.sqrt(d_k)  
    weights = softmax(scores, axis=-1)
    output = weights @ V  
    projected = np.einsum('hsv,hvd->hsd', output, W_o)  
    return projected.transpose(1, 0, 2).reshape(seq_len, -1)

💡 einsum — мощный инструмент для выражения сложных операций с многомерными массивами. Особенно полезен, когда нужно точно контролировать свёртки и трансформации осей. В задачах NLP и computer vision это буквально незаменимая вещь.

📌 Почему стоит обратить внимание:
— Полная векторизация — минимум циклов, максимум скорости;
— Код ближе к математике, а значит — легче проверять;
— Можно выразить довольно сложные операции с тензорами в одной строке.

📍 Навигация: Вакансии • Задачи • Собесы

🐸 Библиотека дата-сайентиста

#буст