ML без магии: одна Tiny Language Model, одна команда Node.js и все веса на виду

English version | Русская версия

Tokenization → embeddings → causal Transformer → LM head → softmax → loss → backpropagation. Без TensorFlow, PyTorch и скрытого autograd.

Репозиторий: tiny-language-model-neuro-js.

Большинство объяснений language model показывает правильные формулы, но прячет путь между ними внутри фреймворка. Мне хотелось обратного: один маленький сценарий, в котором виден каждый скаляр, а терминал прямо показывает неправильные ответы до обучения и правильные после него.

Теперь в проекте одна команда:

node src/train.js --generalize --adaptive-teach

Нужен Node.js 18.19+. Зависимостей нет.

Реальный фрагмент из logs/training-log.txt: матрицы AFTER и DELTA одного FFN-слоя:

Сначала результат

Модель получает вопросы сразу после случайной инициализации:

BEFORE TRAINING — random, usually wrong answers
> can human read ?
  model:    ? <unk> ...
  expected: human can read.  [WRONG]

> can fish swim ?
  model:    ? <unk> ...
  expected: fish can swim.   [WRONG]

> can cat read ?
  model:    ? <unk> ...
  expected: cat cannot read. [WRONG]

После pre-training, SFT и adaptive SFT та же модель отвечает:

FINAL ANSWERS AFTER ADAPTIVE SFT
> can human read ?
  model:    human can read.  [CORRECT]
> can fish swim ?
  model:    fish can swim.   [CORRECT]
> can bird fly ?
  model:    bird can fly.    [CORRECT]
> can cat read ?
  model:    cat cannot read. [CORRECT]

Rehearsal controls preserved: 14/14.
Stable criterion reached 11 times in a row.

Начальный текст меняется из-за случайных весов. Итоговый критерий строгий: все ответы должны быть правильными, вероятность каждого target-токена должна составлять минимум 95%, а полная проверка обязана пройти более десяти раз подряд.

Что осталось после удаления лишних режимов

Раньше код содержал несколько режимов отладки и обучения. Они помогали во время экспериментов, но закрывали главную мысль. В итоговой версии остался один pipeline:

текст → word tokenization → token IDs
      → token + position embeddings
      → два causal Transformer-блока
         → multi-head self-attention
         → FFN с двумя скрытыми слоями
      → LM head → softmax → вероятности следующего токена
      → cross-entropy → backpropagation → Adam

Теперь train.js читается как одна история, а не как маленький CLI-фреймворк.

Скаляр создаёт вычислительный граф

Каждое число, участвующее в обучении, является объектом Value:

class Value {
  constructor(data, children = [], backward = () => {}) {
    this.data = data;
    this.grad = 0;
    this.children = children;
    this._backward = backward;
  }
}

Для умножения:

y = a × b
dy/da = b
dy/db = a

Операция запоминает локальные производные. backward() строит топологический порядок графа и применяет chain rule от итогового loss назад к embeddings и весам.

Нейрон — буквально объект

Формула нейрона не спрятана за tensor API:

output = activation(sum(input[i] × weight[i]) + bias)

Реализация повторяет формулу:

forward(input) {
  let output = sum(
    input.map((value, i) => value.mul(this.weights[i]))
  );

  if (this.useBias) output = output.add(this.bias);
  if (this.activation === 'relu') return output.relu();
  return output;
}

Linear — просто массив нейронов, получающих один вход. Это медленнее матричного умножения, зато каждый вычислительный шаг можно увидеть.

Embeddings и порядок слов

Каждый token ID выбирает обучаемый вектор:

token representation = tokenEmbedding[id] + positionEmbedding[position]

Сначала embeddings содержат случайные числа. Полезные связи появляются только потому, что градиенты многократно изменяют эти числа в учебных контекстах. У cat, read или cannot нет вручную заданного свойства meaning.

Self-attention без сокращений

Для каждого токена:

Q = X × Wq
K = X × Wk
V = X × Wv

score = dot(Q, K) / sqrt(headSize)
attention = softmax(score)
output = attention × V

Цикл работает только при past <= position. Это causal mask: модель видит текущий токен и историю, но никогда не получает будущий target.

После attention каждый токен проходит через FFN с двумя скрытыми слоями:

dModel → hidden ReLU → hidden ReLU → dModel

LayerNorm и residual paths окружают attention и FFN.

Полный шаг обучения

Самый важный код проекта занимает несколько строк:

function learnOneToken({ model, optimizer, input, targetId }) {
  const loss = model.loss(input, targetId);

  optimizer.zeroGrad();
  loss.backward();
  optimizer.step();

  return loss.data;
}

Loss — обычная next-token cross-entropy:

loss = -log(P(target | предыдущие токены))

Если вероятность правильного токена мала, loss велик. Backpropagation вычисляет dLoss/dWeight, Adam изменяет параметры, а следующий forward pass возвращает уже другое распределение.

Этап 1: pre-training

Маленький мир содержит 14 связей способностей:

human can read .
fish can swim .
bird can fly .
dog cannot read .

Связь cat + read намеренно отсутствует. Pre-training выбирает позиции из этого текста и учится предсказывать следующий токен.

Этап 2: SFT

Те же связи превращаются в 42 пары prompt-answer:

can fish swim ?
is fish able to swim ?
does fish know how to swim ?

В SFT loss участвуют только токены ответа. На каждой эпохе реализация проходит каждую пару и каждую позицию ответа, поэтому цикл остаётся детерминированным и понятным.

Этап 3: adaptive SFT

Недостающий ответ представлен только target-токенами:

['cat', 'cannot', 'read', '.']

Эти targets получают шесть вариантов вопроса. Это прямое supervision: модель не открыла зоологический факт самостоятельно. Учитель передал факт через loss, а backpropagation распределил информацию по embeddings, attention, FFN, LayerNorm и LM head.

Почему не остановиться после одного правильного ответа? Потому что одна генерация может быть неустойчивой. Цикл продолжается, пока вероятность каждого target-токена не превысит 95%, а вся проверка не завершится успешно 11 раз подряд.

Catastrophic forgetting и rehearsal

Первая реализация обучала только шесть новых вопросов про кошку. Новый ответ она запомнила, но прежнее поведение разрушилось:

can human read ? → cat cannot read.
can fish swim ?  → cat cannot read.

Это catastrophic forgetting в миниатюре. Исправление — rehearsal: adaptive-эпохи повторяют ещё и 14 старых примеров can ... ?. Финальный критерий проверяет одновременно новые и старые ответы, поэтому модель не может закончить обучение, просто заменив всё одним ответом.

Лог создаётся всегда

Команда автоматически записывает:

logs/training-log.txt

Это последовательная ASCII-схема, а не сырой JSON-дамп. Сначала в ней идут все
события forward/loss/backward/update, затем полные матрицы в трёх контрольных
точках:

начальные случайные матрицы
        |
        v
матрицы после pre-training + SFT
        |
        v
финальные матрицы после adaptive SFT

Для каждого перехода лог печатает AFTER-матрицу и точную DELTA-матрицу.
Строки Linear-слоёв подписаны как neuron[n], столбцы — как weight[n],
а bias показан рядом со своим нейроном. Отдельно отмечается самое сильное
изменение в формате слой / нейрон / вес: before -> after -> delta.

Насколько это близко к production LLM

Архитектура и правило обучения настоящие, масштаб намеренно маленький.

Эта модель	Production model
24 word-токена	Большой subword/byte-level словарь
2 160 параметров	Миллионы или миллиарды
Два Transformer-блока	Десятки или сотни
Скалярный JavaScript-граф	Батчевый тензорный граф на ускорителях
Маленький структурированный корпус	Огромные подготовленные datasets
Узкое обученное поведение	Широкий язык и reasoning

Проект не конкурирует с GPT. Это causal language model, уменьшенная до размера, при котором весь путь помещается в одном репозитории и одной мысленной модели:

токен → embedding → attention → FFN → probability
      → loss → gradient → новый вес → изменившийся ответ

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

Репозиторий: tiny-language-model-neuro-js.

Автор: Максим Секретов.