Как сократить размер BERT без значительной потери в точности

BERT — это одна из state-of-the-art моделей в обработке естественного языка. Несмотря на качество модели, обучение BERT является ресурсоемким как по памяти, так и по скорости обучения. Исследователи из RASA протестировали техники масштабирования нейросетей на BERT. С помощью квантизации и TFLite размер BERT сократился в 4 раза.

Техники масштабирования нейросетей

Квантизация

Квантизация означает уменьшение численной точности весов модели. Один из популярных методов — k-means квантизация. Имея веса модели в матрице W с десятичными числами, веса кластеризуются с помощью k-means в N кластеров. Затем матрица W трансформируется в матрицу целых чисел от 1 до N, каждое из которых является указателем к центру кластера. Так можно сжать каждый элемент изначальной матрицы из 32-битного десятичного числа в log(N)-битные целые числа.

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

Квантовать веса;
Посчитать функцию потерь в квантованной нейросети;
Рассчитать градиенты по функции потерь с учетом не квантованных весов;
Обновить не квантованные веса

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

Прунинг

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

Дистилляция знаний

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

Ускорение BERT

Цель исследования была в том, чтобы ускорить предсказание BERT, чтобы использовать нейросеть для классификации именованных сущностей. На момент экспериментов классификатор с BERT требовал ~120ms на CPU, чтобы обработать одно сообщение. Другие классификаторы по скорости были на ~100x быстрее.

Квантизация с TFLite

В качестве первого эксперимента исследователи протестировали 8-битную квантизацию с TensorFlow Lite. Результаты были следующими:

Квантизация не повлияла на точность модели: F1 метрика BERT модели сократилась на 0.3%;
Размер BERT сократился в 4 раза: веса конвертировались из 32-битных чисел в 8-битные;
Скорость работы BERT на TensorFlow Lite увеличилась в 5 раз на CPU