Self-attention

Self-attention — механизм, лежащий в основе современных архитектур трансформеров, таких как BERT, GPT и других. В контексте TensorFlow.js он позволяет эффективно обрабатывать последовательности данных, выявлять зависимости между элементами и формировать представления с учётом контекста.

Тензоры и подготовка данных

Основной строительный блок в TensorFlow.js — tf.Tensor. Для работы с self-attention требуется подготовить входные последовательности в виде тензоров:

const tf = require('@tensorflow/tfjs');

// Пример случайной последовательности длиной 5 и размерностью признаков 8
const sequenceLength = 5;
const featureDim = 8;
const input = tf.randomNormal([1, sequenceLength, featureDim]);

Здесь [1, 5, 8] означает [batch_size, sequence_length, feature_dimension].

Проекция ключей, запросов и значений

Self-attention основывается на трёх компонентах:

  • Query (Q) — запросы
  • Key (K) — ключи
  • Value (V) — значения

Для их формирования используются плотные слои (Dense Layers):

const queryLayer = tf.layers.dense({ units: featureDim });
const keyLayer = tf.layers.dense({ units: featureDim });
const valueLayer = tf.layers.dense({ units: featureDim });

const Q = queryLayer.apply(input);
const K = keyLayer.apply(input);
const V = valueLayer.apply(input);

Каждый слой проецирует входные данные в новое пространство признаков для вычисления внимания.

Вычисление матрицы внимания

Классический подход — Scaled Dot-Product Attention. Основные шаги:

  1. Умножение запросов на транспонированные ключи.
  2. Масштабирование на корень из размерности ключа.
  3. Применение softmax для нормализации.
const matmulQK = tf.matMul(Q, K, false, true);
const dk = tf.scalar(Math.sqrt(featureDim));
const scaledScores = tf.div(matmulQK, dk);
const attentionWeights = tf.softmax(scaledScores);
  • tf.matMul(Q, K, false, true) вычисляет скалярное произведение каждого запроса с ключами.
  • Масштабирование предотвращает чрезмерно большие значения перед softmax.
  • Softmax обеспечивает нормированные веса, отражающие важность каждого элемента последовательности.

Применение весов внимания к значениям

Для получения окончательного представления последовательности необходимо умножить веса внимания на значения:

const output = tf.matMul(attentionWeights, V);

output имеет ту же форму, что и входная последовательность, но каждый элемент теперь учитывает контекст всей последовательности.

Многофоновое (Multi-Head) внимание

Для повышения выразительности модели используется несколько параллельных голов внимания. Каждая голова проецирует данные в отдельное подпространство признаков и вычисляет своё внимание:

const numHeads = 4;
const headDim = featureDim / numHeads;

function splitHeads(x) {
  return tf.reshape(x, [x.shape[0], x.shape[1], numHeads, headDim])
           .transpose([0, 2, 1, 3]); // [batch, heads, seq_len, headDim]
}

const Q_heads = splitHeads(Q);
const K_heads = splitHeads(K);
const V_heads = splitHeads(V);

// Вычисление внимания для каждой головы
const attention_heads = tf.matMul(Q_heads, K_heads, false, true)
                          .div(tf.scalar(Math.sqrt(headDim)))
                          .softmax()
                          .matMul(V_heads);

// Объединение голов обратно
const concatHeads = attention_heads.transpose([0, 2, 1, 3])
                                    .reshape([1, sequenceLength, featureDim]);

Многофоновое внимание позволяет модели учитывать разнообразные зависимости и связи между элементами последовательности, повышая качество извлечения контекста.

Маски внимания

При работе с последовательностями разной длины или при предотвращении «заглядывания вперёд» в задачах генерации текста используется маска:

const mask = tf.tensor2d([[0, 0, 1, 1, 1]]); // 0 — учесть, 1 — замаскировать
const negativeInf = tf.scalar(-1e9);
const maskedScores = scaledScores.add(mask.mul(negativeInf));
const maskedAttention = tf.softmax(maskedScores);

Маска корректирует значения перед softmax, чтобы исключить нежелательные элементы.

Интеграция с TensorFlow.js Layers API

Self-attention можно встроить в модель как пользовательский слой:

class SelfAttention extends tf.layers.Layer {
  constructor(config) {
    super(config);
    this.units = config.units;
  }

  build(inputShape) {
    this.queryLayer = this.addWeight('query', [inputShape[2], this.units]);
    this.keyLayer = this.addWeight('key', [inputShape[2], this.units]);
    this.valueLayer = this.addWeight('value', [inputShape[2], this.units]);
  }

  call(input) {
    const Q = tf.matMul(input, this.queryLayer.read());
    const K = tf.matMul(input, this.keyLayer.read());
    const V = tf.matMul(input, this.valueLayer.read());

    const scores = tf.matMul(Q, K, false, true)
                     .div(tf.scalar(Math.sqrt(this.units)));
    const weights = tf.softmax(scores);
    return tf.matMul(weights, V);
  }

  computeOutputShape(inputShape) {
    return inputShape;
  }
}

Это позволяет использовать self-attention в составе сложных моделей и обучать его совместно с другими слоями.

Практические советы

  • Для больших последовательностей использовать масштабирование и маскирование обязательно, чтобы избежать переполнения и неправильного внимания.
  • Регуляризация через dropout в слоях внимания помогает бороться с переобучением.
  • Вычисление многоголовного внимания в TensorFlow.js выгодно распараллеливать через reshape и transpose, чтобы уменьшить накладные расходы на циклы JavaScript.

Self-attention в TensorFlow.js обеспечивает гибкий и высокоэффективный способ работы с последовательными данными, позволяя моделям учитывать сложные зависимости внутри последовательностей.

Оглавление