Как сделать трансформер чат-бот на Trax?

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

Сравнительно простого чат-бота можно построить на базе языковой модели, которая умеет прогнозировать следующее слово по предыдущим, и которую несложно сделать, используя Трансформер-декодер по аналогии с GPT. В этом случае диалог формируется как связный неструктурированный текст. Чтобы превратить этот текст в чат, нужно вмешиваться в процесс генерации, добавляя реплики пользователя. Но обо всём по порядку.

Для модели чат-бота я решил попробовать новую архитектуру — Reformer. Это трансформер, который может работать с длинными текстами размером с хорошую книгу буквально на одном ускорителе. Reformer кардинально снижает используемый объем памяти. Это достигается за счет двух вещей. Во-первых, у Reformer более эффективный в смысле памяти механизм внимания. Во-вторых, реверсивная схема вычислений, позволяющая отказаться от хранения значений активации для расчета градиента. Детали здесь на английском, а здесь перевод заметки из Google AI Blog.

Мотивация: не то что бы у меня были какие-то гигантские диалоги, но очень хотелось посмотреть как Reformer ведет себя на Google Colab и как учится в сравнении с Transformer.
Оказалось, не зря.

Данные

Нужный набор данных я скачал на Яндекс.Толока. Он называется «Toloka Persona Chat Rus» и его можно использовать в некоммерческих целях с упоминанием источника. Упоминаю: Ребята из Яндекс.Толока — вы большие молодцы!

Набор содержит 10000 русскоязычных диалогов на общие темы с возможностью фильтрации по профилю пользователя. Можно было бы например отфильтровать диалоги, где один из пользователей женщина, но я посчитал, что их и так не очень много. Диалоги представлены в виде HTML-текста, пришлось потратить некоторое время на чистку.

Для обучения модели требуется неструктурированный текст, примерно такой:
'1: Привет) расскажи о себе 2: Привет) под вкусный кофеек настроение поболтать появилось ) 1: Что читаешь? Мне нравится классика Я тоже люблю пообщаться 2: Люблю животных, просто обожаю, как и свою работу) Я фантастику люблю 1: А я выращиваю фиалки И веду здоровый и активный образ жизни! 2: Ух ты, интересно. 1: Ты случайно не принц на белом коне? Я его очень жду.. 2: А у меня из хобби каждую неделю тусить с моим лучшим другом) STOP'

Два момента:

1. Идентификаторы «Пользователь 1:» и «Пользователь 2:» в исходном тексте пришлось сократить до «1:» и «2:», чтобы они не разбивались на два токена.

2. В конец каждого диалога я добавил слово STOP, которое кодируется одним токеном. Это нужно для остановки декодера при генерации диалога.

Справедливо решив, что для модели с несколькими десятками миллионов параметров десяти тысяч диалогов будет маловато, я сделал еще 36 тысяч, просто «откусывая» от начала диалогов по четному количеству реплик так, чтобы оставалось не меньше восьми. В общем, и этого не очень-то много, в связи с этим вопрос: где взять еще диалогов на русском?

Процесс подготовки данных представлен на схеме:

В качестве модели для сегментации текста, как и в эксперименте с саммаризатором, я использовал Byte Pair Encoding (BPE) из библиотеки sentencepiece. Размер словаря 10k токенов.

Размеры корзин для формирования пакетов фиксированной длинны выбирал по гистограмме. Решил отказаться от текстов короче шестидесяти четырех токенов.

Модель

Модель Reformer создается несколькими строчками кода:

def ReformerLM(vocab_size=vocab_size, n_layers=6, mode='train', 
               attention_type=tl.SelfAttention):

    model = trax.models.ReformerLM( 
        vocab_size=vocab_size,
        n_layers=n_layers,
        mode = mode,
        attention_type=attention_type
    )
    return model
  
ChatModel = ReformerLM()

У Reformer есть два собственных варианта реализации механизма внимания: Self Attention и LSH Attention. Первый экономит память, рассчитывая только отдельные сегменты матрицы весов внимания, хотя это по прежнему dot-product attention. Второй оценивает сходство векторов через локально-чувствительное хеширование (locality-sensitive hashing). LSH я еще не испытал, по идее должно работать быстрее, но нужно подбирать параметры хеширования.

За исключением размера словаря параметры модели взяты по умолчанию:
Шесть декодеров, в каждом восемь головок
Размер словаря: 10000
Размер векторного пространства: 512
Размер слоя прямого распространения: 2048
Вероятность выключения нейронов из обучения (dropout): 0.1

Схему модели можно вывести командой print(model).

Картинка в стиле «длиннокот» для одного декодер-блока

Обучение

В моей модели 35 млн. параметров, обучал сериями по 20 тысяч шагов. Каждая серия чуть больше двух часов. Всего восемь серий или 160 тысяч итераций. Загрузка весов для модели после очередной серии является стрессом, качество в этот момент нестабильно. Заметил это только когда «прикрутил» TensorBoard. Оказалось совсем несложно, формат логов стандартный.

TensorBoard в colab-ноутбуке:

%load_ext tensorboard
%tensorboard --logdir ../root/model/ # folder where train and eval logs...

Кривые обучения для третьей и четвертой серии. Виден «стресс» на 40k.

До 100k с шагом обучения 2e-4 модель учится сравнительно ровно. Потом проявляются переобучение и зависания в локальных минимумах. Нужно снижать шаг обучения и увеличивать dropout. Не рекомендую dropout больше 0.2. Обучение сильно замедляется, а качество откатывается на пару серий назад. Кривые для шестой и седьмой серии:

За восемь серий кросс-энтропийная ошибка (cross entropy loss) упала с 10 до 1.6-1.8, доля правильных ответов (accuracy) выросла в среднем до 60%.

Чат-бот

Для чат-бота нужен интерактивный декодер (не путать с декодер-блоком Reformer). Он должен принимать реплики пользователя и добавлять их в конец последовательности, на основе которой модель определяет следующий токен. В основе декодера — авторегрессионный генератор из библиотеки trax. Авторегрессионный означает, что следующий токен генерируется на основе текущего и всей предыдущей истории, а затем сам становится текущим. Чтобы очистить историю нужно установить исходное состояние модели, которое можно считать в переменную после инициализации и загрузки весов. Делается это простой инструкцией: STARTING_STATE = model.state и также потом устанавливается: model.state = STARTING_STATE.

def output_gen(model, start_sequence, temperature, accelerate):
        
    # Add batch dimension to array
    input_batch = np.expand_dims(start_sequence, axis=0)
    
    gen = trax.supervised.decoding.autoregressive_sample_stream( 
        model,
        inputs=input_batch,
        temperature=temperature,
        accelerate=accelerate
    )

    return gen

Я сделал два декодера, бот — первый, пользователь — второй, и наоборот. Во всех примерах жадный декодер (температура = 0, см. спойлер "Детали декодера"). Если температура нулевая и вы пишите одни и те же реплики, диалог будет одним и тем же.

Пример почти идеального диалога:
(бот — первый, реплики выделены):
1: Привет
2: Привет
1: Как дела?
2: Отлично. А у вас?
1: Тоже хорошо Чем занимаешься?
2: Я фрилансер
1: Я учитель младших классов
2: Ух ты. Вам нравятся дети?
1: Да, у меня трое детей
2: Наверное нелегко с ними?
1: Да, мы с ними часто ходим в походы
2: Ого. Вы любите путешествовать?
1: Да, люблю
2: А я хочу в Японию. Это моя мечта
1: Что ты любишь?
2: Музыку и красное вино. А вы?
1: Мне нравится
2: Ладно. Мне пора бежать.

Скринкаст. Ускорил в два раза.

В коде есть еще full text decoder. это генератор текста без пользователя. Он нужен, чтобы быстро оценить качество этапа.

Для декодеров использовал фрагменты кода из неоцениваемой части задания на курсе "Natural Language Processing with Attention Models". Курс великолепен. Рекомендую!

Ссылки

• Мой репозитарий с кодом эксперимента.

• Репозитарий trax

• Статья на arxiv: "Reformer: The Efficient Transformer", вышла в начале 2020 года .