Reformer на TRAX? / Habr

Что такое Reformer и почему он круче Transformer’a?

Давайте предварительно начнем с того, что же такой Reformer и почему благодаря ему мы можем рассчитывать на расширение контекстов вплоть до десятков тысяч слов.

В классической архитектуре Transformer механизм внимания работает со сложностью, которая масштабируется квадратично с увеличением длины последовательности.

Это происходит потому, что каждый токен в последовательности должен вычислять оценки внимания со всеми другими токенами, что приводит к плотной матрице внимания, размер которой растет с квадратом длины последовательности – мрак для вычислительных способностей наших TPU и GPU.

Reformer решает эту проблему с помощью locality-sensitive hashing. LSH преобразует способ вычисления внимания, локализуя/приближая полный механизм внимания. Вместо вычисления оценок внимания между всеми парами токенов, LSH хеширует токены последовательности в корзины таким образом, что похожие токены с большей вероятностью попадают в одну и ту же корзину.

Как это работает? Допустим, каждая картинка превращается в набор чисел (вектор), и мы хотим найти картинки с похожими наборами чисел. LSH использует специальные математические функции, называемые хэш-функциями, чтобы превратить каждый набор чисел в короткую строку битов (хэш-код). Ключевая идея LSH состоит в том, что похожие наборы чисел должны давать похожие хэш-коды.

Теперь у нас есть много хэш-кодов для всех наших картинок. Мы можем разделить все хэш-коды на корзины и положить хэш-коды, которые похожи друг на друга, в одну корзину. Таким образом, картинки, которые были похожи друг на друга, теперь находятся в одной группе – корзине.

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

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

По итогу по мере роста глубины нейросети – мы получаем вновь высоченные требования к вычислениям.

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

Эти два ключевых нововведения в Reformer: использование LSH для эффективного вычисления внимания и обратимых остаточных слоев для оптимизации памяти — позволяют обрабатывать гораздо более длинные последовательности с существенно меньшими затратами ресурсов по сравнению с классическим Transformer. Теперь перед нами статьи из Википедий и целые книги…

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

Хотя и классический Transformer, и Reformer стремятся эффективно моделировать зависимости в последовательных данных, архитектурные инновации Reformer устраняют ограничения масштабируемости, присущие классическому Трансформеру.

Reformer достигает большей эффективности и масштабируемости, делая его мощным инструментом для обработки длинных последовательностей в обработке естественного языка, геномике, альбомов рэп-исполнителей… Ну да ладно…

В заключение, Reformer представляет собой глубокий сдвиг в дизайне нейронных сетей для моделирования последовательностей, устраняя основные ограничения традиционных Transformers. Короче говоря, мы получили действительно “эффективный” трансформер, который не требует от нас насильного изъятия рендер-ферм бедных трехмерщиков и активации ядерного реактора под мегаTPU.

Но гении нашего мира догадались использовать JAX и Reformer…

Библиотека для эффективного программирования?

Вот знаете, разрабатывали еще фреймворки для эффективной работы. JAX (Just Another eXtensor library) – библиотека для вычислений на графических процессорах (GPU) и наших любимых процессорах (TPU), автоматического дифференцирования и других вычислений в Python.

Одна из ключевых особенностей JAX – автоматическое дифференцирование. В отличие от методов разделения вручную или символьного дифференцирования, автоматическое разделение позволяет вычислять производные программно, используя компьютерные алгоритмы.

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

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

JAX также внутренне использует XLA (Accelerated Linear Algebra, специальный компилятор, для оптимизации и ускорения выполнения вычислений. Тут неплохая статья с объяснением, что же это за зверь такой…

Это позволяет JAX достигать высокой производительности на GPU и TPU, делая его привлекательным выбором для задач, требующих высокой вычислительной мощности.

Да более того, JAX поощряет функциональный стиль программирования, в котором функции остаются неизменяемыми, и состояние передается через аргументы – получается чистый и распределенного кода.

Но вот появилась одна из задач реализовать реформер на библиотеке, заточенной под работу в контексте того самого JAX. Для этого нужно было как-то объединить и скоординировать JAX и классический TensorFlow – так появился TRAX от ресерчеров из Гугла, который работает поверх библиотек.

Она разработана командой Google Brain и спроектирована таким образом, чтобы быть эффективной и простой в использовании, учитывая как исследовательские, так и практические потребности. Благодаря своей модульности

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

Как работать с библиотекой?

Она предлагает различные заранее определенные слои, активации и функции потерь, позволяя пользователям создавать пользовательские модели, настроенные под их конкретные потребности. У Trax есть своя коллекция предварительно обученных моделей: Трансформер и ResNet, которые легко можно донастроить или использовать для задач вывода.

Блочное программирование в Trax реализуется с помощью композиции слоев (layer compositions). Короче говоря, она работает по принципу блочного конструирования.

Слои – отдельные функции, принимающие на вход некоторые данные и преобразующие их в соответствии с определенными правилами. Например, слой Dense выполняет операцию умножения входных данных на матрицу весов и добавления смещения.

Для создания более сложных моделей в Trax можно объединять слои вместе с помощью функции Serial, которая последовательно применяет слои друг за другом, создавая цепочку слоев, где вывод одного слоя является входом для следующего.

Также можно комбинировать несколько слоев параллельно с помощью функции Parallel, которая применяет указанные слои к одним и тем же данным и объединяет их выводы.

Для повторяющихся шаблонов использования слоев Trax предоставляет функцию Branch, которая позволяет повторять один и тот же слой несколько раз внутри композиции.

Вот пример создания самого обычного слоя в Trax.

import trax
from trax import layers

# Создание слоя полносвязного перцептрона
dense_layer = layers.Dense(n_units=64)

# Применение слоя к входным данным
input_data = trax.random.uniform(shape=(1, 128))
output = dense_layer(input_data)

print(output.shape)  # Результат — форма выходных данных

Когда слои собраны в цепочки, они образуют модели. Вот пример конструирования и компиляции уже полноразмерной модели в Trax:

import trax
from trax import layers

# Создание модели с двумя слоями Dense
model = trax.models.MLP([layers.Dense(128), layers.Dense(64)])

# Компиляция модели с указанием функции потерь и оптимизатора
model.compile(loss_fn=trax.layers.CrossEntropyLoss(), optimizer=trax.optimizers.Adam(0.01))

# Вывод структуры модели
print(model)

Этот пример демонстрирует создание и компиляцию модели с двумя полносвязными слоями. Модель готова к обучению с использованием функции потерь CrossEntropyLoss и оптимизатора Adam с коэффициентом обучения 0.01.

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

# Загрузка данных для обучения и тестирования
train_generator = trax.data.inputs.NumpyReader(input_data=train_inputs, input_steams=train_targets)
eval_generator = trax.data.inputs.NumpyReader(input_data=eval_inputs, input_streams=eval_targets)

# Обучение модели
model.fit(train_generator, epochs=10, eval_steps=eval_steps, eval_freq=1)

# Оценка модели
eval_loss = model.evaluate(eval_generator)
print("Evaluation Loss:", eval_loss)

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

Вот пример загрузки предварительно обученных весов в модель Trax:

import trax
from trax import layers

# Создание модели
model = trax.models.TransformerLM(  # Например, модель TransformerLM
    d_model=512,
    d_ff=2048,
    n_layers=6,
    n_heads=8,
    vocab_size=32000,
    max_len=4096,
    mode='train'
)

# Загрузка предварительно обученных весов
pretrained_weights = '/path/to/pretrained_model_weights.npz'
model.init_from_file(pretrained_weights)

# Просмотр структуры модели
print(model)

В этом примере модель TransformerLM создается с заданными параметрами, а затем загружаются предварительно обученные веса из файла.

Это позволяет использовать знания, накопленные предварительно обученной моделью, для последующей тонкой настройки на специфические задачи.

После загрузки предварительно обученных весов, модель может быть адаптирована к конкретной задаче путем тонкой настройки (fine-tuning) ее весов на новых данных.

Вот пример адаптации предварительно обученной модели для классификации текста:


import trax
from trax import layers

# Создание модели
model = trax.models.TransformerLM(  # Например, модель TransformerLM
    d_model=512,
    d_ff=2048,
    n_layers=6,
    n_heads=8,
    vocab_size=32000,
    max_len=4096,
    mode='train'
)

# Загрузка предварительно обученных весов
pretrained_weights = '/path/to/pretrained_model_weights.npz'
model.init_from_file(pretrained_weights)

# Создание адаптированной версии модели для классификации текста
classification_model = trax.models.TransformerClassifier(
    body=model,
    n_classes=2  # Например, классификация на 2 класса
)

# Просмотр структуры адаптированной модели
print(classification_model)

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

Одно из преимуществ библиотеки – ее интеграция с TensorFlow.

Обслуживание моделей с использованием TensorFlow Serving включает развертывание обученных моделей машинного обучения в производственных средах, позволяя им предоставлять прогнозы или выводы через протоколы HTTP или gRPC. TensorFlow Serving предоставляет гибкую и высокопроизводительную инфраструктуру для обслуживания моделей машинного обучения в различных сценариях развертывания.

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

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

Экспорт модели: Перед тем, как обслуживать модель с помощью TensorFlow Serving, необходимо экспортировать ее в формат, совместимый с TensorFlow Serving. Обычно это включает сохранение обученной модели в формате SavedModel. Ниже приведен пример базового способа экспорта модели в TensorFlow:

import tensorflow as tf

# Предположим, что 'model' - ваша обученная модель TensorFlow
model.save('/путь/к/каталогу/экспорта', save_format='tf')

Установка TensorFlow Serving: TensorFlow Serving должен быть установлен на сервере, на котором вы планируете развернуть вашу модель. Его можно установить с помощью Docker, менеджеров пакетов или собрать из исходного кода.

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

Запуск сервера TensorFlow Serving: После настройки запустите сервер TensorFlow Serving. Это запустит сервер и загрузит вашу экспортированную модель в память, готовую к обслуживанию предсказаний.

Отправка запросов на сервер: При запущенном сервере TensorFlow Serving можно отправлять запросы на предсказания. C использованием HTTP REST API или gRPC. Клиенты могут отправлять входные данные на сервер, который обработает их с использованием загруженной модели и вернет предсказания.

Простой пример запуска сервера TensorFlow Serving с Docker:

docker run -p 8501:8501 --name=tfserving_model \
--mount type=bind,source=/путь/к/экспортированной_модели,target=/models/model_name \
-e MODEL_NAME=model_name -t tensorflow/serving

Эта команда запускает сервер TensorFlow Serving с использованием Docker, открывая его на порту 8501 и загружая экспортированную модель из каталога /путь/к/экспортированной_модели.

Reformer в коллабе с Trax?

Прежде всего, необходимо установить библиотеку TRAX. Это можно сделать с помощью команды pip install trax. После успешной установки можно приступать к созданию модели. Для этого нужно импортировать необходимые модули и задать конфигурацию модели.

Первый щаг - импорт нужных модулей:

import trax
from trax import layers as tl
from trax.supervised import training

Следующим этапом является определение гиперпараметров модели, таких как размер словаря (vocab_size), размер эмбеддингов (d_model), количество слоев (n_layers), количество голов в механизме внимания (n_heads) и другие параметры, которые определяют структуру и поведение модели. Пример задания гиперпараметров может выглядеть следующим образом:

vocab_size = 32000
d_model = 512
n_layers = 6
n_heads = 8
max_len = 2048

После задания гиперпараметров необходимо определить саму модель. В TRAX это делается с помощью последовательного создания слоев. Основной блок Реформера:

def ReformerLM(vocab_size=vocab_size, d_model=d_model, n_layers=n_layers, n_heads=n_heads, max_len=max_len):
    return trax.models.reformer.ReformerLM(
        vocab_size=vocab_size,
        d_model=d_model,
        n_layers=n_layers,
        n_heads=n_heads,
        max_len=max_len
    )

Здесь ReformerLM – это функция, создающая экземпляр модели ReformerLM с заданными параметрами.

Следующим шагом является подготовка данных для обучения. В TRAX имеется удобный модуль для работы с датасетами.

train_stream = trax.data.TFDS('wmt32k', keys=('inputs', 'targets'), train=True)
eval_stream = trax.data.TFDS('wmt32k', keys=('inputs', 'targets'), train=False)

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

Для обучения модели необходимо создать обучающий цикл. В TRAX это делается с помощью класса training.TrainTask.

train_task = training.TrainTask(
    labeled_data=train_stream,
    loss_layer=tl.CrossEntropyLoss(),
    optimizer=trax.optimizers.Adam(),
    lr_schedule=trax.lr.warmup_and_rsqrt_decay(4000, 0.01),
    n_steps_per_checkpoint=500
)

Аналогичным образом создается и валидационная задача:

eval_task = training.EvalTask(
    labeled_data=eval_stream,
    metrics=[tl.CrossEntropyLoss(), tl.Accuracy()],
    n_eval_batches=10
)

После создания тренировочной и валидационной задач можно определить обучение, объединив их в обучающий процесс. Используем training.Loop:

training_loop = training.Loop(
    ReformerLM(),
    train_task,
    eval_tasks=[eval_task],
    output_dir='/content/reformer_model/'
)

Обучение модели запускается командой training_loop.run(n_steps=10000), где n_steps – это количество шагов, которые модель должна пройти в процессе обучения.

И дополнительно рекомендуем почитать:

https://github.com/google/jax

https://ru.wikipedia.org/wiki/Google_JAX

https://habr.com/ru/articles/522622/

https://paperswithcode.com/method/reformer

https://arxiv.org/abs/2001.04451v2

https://arxiv.org/pdf/2001.04451v2

https://research.google/blog/reformer-the-efficient-transformer/

https://github.com/google/trax

https://trax-ml.readthedocs.io/en/latest/notebooks/trax_intro.htm

Reformer на TRAX?

{{ titleHtml }}

{{ titleHtml }}