Построение упрощённого трансформера в TensorFlow.js

TensorFlow.js — это библиотека для машинного обучения на JavaScript, которая позволяет создавать, обучать и использовать модели прямо в браузере или на сервере с Node.js. Основной элемент работы с библиотекой — это тензоры (tf.Tensor), которые являются многомерными массивами чисел, поддерживающими автоматическое дифференцирование.

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

Тензоры и операции

Тензоры создаются из массивов Jav * aScript:

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

const a = tf.tensor([1, 2, 3, 4]);
const b = tf.tensor2d([[1, 2], [3, 4]]);

Ключевые операции включают matMul для матричного умножения, add, sub, mul и div для арифметических операций, а также функции активации (relu, softmax).

Позиционные кодировки

Поскольку трансформеры не используют рекуррентные сети, они требуют позиционной информации. Позиционная кодировка добавляется к эмбеддингам входных токенов:

function positionalEncoding(maxLen, dModel) {
  const pos = tf.range(0, maxLen);
  const i = tf.range(0, dModel, 2);
  const angleRates = tf.div(pos.expandDims(1), tf.pow(10000, tf.div(i, dModel)));
  
  const sinEncoding = tf.sin(angleRates);
  const cosEncoding = tf.cos(angleRates);
  
  return tf.concat([sinEncoding, cosEncoding], -1);
}

Ключевой момент: позиционная кодировка позволяет модели различать порядок элементов в последовательности, сохраняя параллельную обработку.

Механизм внимания

Self-Attention — центральная часть трансформера. Он вычисляет вес каждого элемента последовательности относительно всех остальных:

  1. Создаются запросы (Q), ключи (K) и значения (V) через линейные слои:
function denseLayer(input, units) {
  return tf.layers.dense({ units }).apply(input);
}

const Q = denseLayer(inputTensor, dModel);
const K = denseLayer(inputTensor, dModel);
const V = denseLayer(inputTensor, dModel);
  1. Вычисляется скалярное произведение и нормализуется через softmax:
function scaledDotProductAttention(Q, K, V) {
  const matmulQK = tf.matMul(Q, K, false, true);
  const dk = tf.scalar(K.shape[K.shape.length - 1], 'float32');
  const scaledScores = tf.div(matmulQK, tf.sqrt(dk));
  const attentionWeights = tf.softmax(scaledScores);
  return tf.matMul(attentionWeights, V);
}

Особенности реализации:

  • Маски используются для предотвращения “заглядывания вперёд” при генерации текста.
  • Подразумевается, что размерности Q, K, V согласованы с dModel.

Многоголовое внимание

Многоголовое внимание позволяет модели одновременно учитывать разные представления последовательности:

function multiHeadAttention(input, numHeads, dModel) {
  const depth = dModel / numHeads;
  const heads = [];

  for (let i = 0; i < numHeads; i++) {
    const Q = denseLayer(input, depth);
    const K = denseLayer(input, depth);
    const V = denseLayer(input, depth);
    const head = scaledDotProductAttention(Q, K, V);
    heads.push(head);
  }

  return tf.concat(heads, -1);
}

Слои нормализации и остаточные соединения

Трансформеры используют layer normalization и residual connections для стабильного обучения:

function addResidualNorm(x, sublayer) {
  const out = tf.add(x, sublayer(x));
  return tf.layers.layerNormalization().apply(out);
}

Это помогает сохранять градиенты и ускоряет сходимость.

Построение энкодера

Энкодер состоит из N одинаковых блоков, каждый из которых включает:

  1. Многоголовое внимание
  2. Слой нормализации с остаточным соединением
  3. Feed-forward сеть
  4. Второй слой нормализации с остаточным соединением

Пример одного блока:

function encoderBlock(input, numHeads, dModel, dff) {
  const attnOutput = addResidualNorm(input, x => multiHeadAttention(x, numHeads, dModel));
  const ffnOutput = addResidualNorm(attnOutput, x => tf.layers.dense({ units: dff, activation: 'relu' }).apply(x));
  return ffnOutput;
}

Блоки энкодера можно стекать, формируя полную модель:

let encoderOutput = inputEmbedding;
for (let i = 0; i < numBlocks; i++) {
  encoderOutput = encoderBlock(encoderOutput, numHeads, dModel, dff);
}

Feed-Forward слой

Feed-Forward слой применяет нелинейное преобразование к каждому элементу последовательности:

const ffn = tf.sequential();
ffn.add(tf.layers.dense({ units: dff, activation: 'relu' }));
ffn.add(tf.layers.dense({ units: dModel }));

Он повышает способность модели захватывать сложные зависимости.

Обучение и оптимизация

Для обучения трансформера применяются стандартные методы TensorFlow.js:

model.compile({
  optimizer: tf.train.adam(0.001),
  loss: tf.losses.softmaxCrossEntropy,
  metrics: ['accuracy']
});

Важно учитывать:

  • Пакетная обработка последовательностей одинаковой длины.
  • Маскирование паддингов для корректной работы attention.
  • Снижение learning rate с использованием warm-up шагов для стабильности.

Инференс и генерация последовательностей

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

function generateNextToken(inputSeq) {
  const logits = model.predict(inputSeq);
  const nextToken = tf.argMax(logits.slice([-1, 1]), -1);
  return nextToken;
}

Механизм attention и позиционные кодировки обеспечивают правильный порядок и согласованность генерируемого текста.

Вывод

TensorFlow.js позволяет полностью реализовать упрощённый трансформер на JavaScript, используя стандартные компоненты: тензоры, слои Dense, LayerNormalization, softmax и матричные операции. Правильная организация внимания, позиционной информации и feed-forward слоев обеспечивает базовую функциональность современных трансформеров, пригодную для экспериментов и обучения небольших моделей прямо в браузере.

Оглавление