Сложный

11 мин

0

Уже прошло два года с тех пор, как мы предложили схему 4.6-битного квантования и рассказали про нее, в том числе и на Хабре: раз и два. Вспомним, что при 4.6-битном квантовании веса и входы слоя принимают такие целые значения, что их попарные произведения помещаются в знаковый 8-битный тип данных. Такая схема позволила нам вычислять нейронные сети на процессорах мобильных устройств быстрее, чем в 8-битном формате, и точнее, чем в 4-битном, потому что уровней квантования больше.

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

Что мешало практическому применению 4.6-битных сетей?

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

Часть этих требований уже закрывает схема 4.6-битного квантования сама по себе. Она позволяет компактно хранить веса и быстро выполнять низкоразрядные матричные операции (свёртки и матричные умножения) на процессоре.

С оперативной памятью ситуация менее однозначная. На вход матричной операции приходят 4.6-битные числа, которые мы храним в 8-битных значениях в оперативной памяти для ускорения вычислений. Такой компромисс приемлем для небольших распознающих сетей, работающих на мобильном устройстве. Однако выход матричной операции — это набор 32-битных целых чисел. Этот результат нужно передать дальше: применить активацию и  подготовить вход для следующего слоя, снова применив 4.6-битное квантование. То есть у нас возникают промежуточные 32-битные карты признаков (до и после активации). А это увеличивает расход оперативной памяти в 4 раза, что может быть крайне нежелательно, например при обработке изображения в высоком разрешении.

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

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

Можно!

HardTanh — наш ответ

В качестве такой активации удобно использовать HardTanh. Это простая ограниченная функция: внутри заданного диапазона она ведёт себя линейно, а всё, что выходит за границы, обрезает до граничных значений.

В простейшем виде:

В более общем случае можно задать ширину диапазона:

По сути такая активация симметрично ограничивает диапазон значений. Если значение лежит внутри диапазона – оно не меняется, если снаружи – ограничивается значениями.

Раньше мы выяснили, что для 4.6-битной сети лучшая точность получается, когда числа уровней квантования для весов и активаций примерно равны. Так что для практического использования выберем схему, где и у весов, и у активаций 23 уровня. Т.е. каждый отдельный вес или вход слоя принимают значения от -11 до +11. Произведения лежат в диапазоне от -121 до 121, что помещается в 8-битный знаковый тип, а значит удовлетворяет условиям 4.6-битного квантования.

Тогда равномерное квантование входа можно представить как 

где s — шаг квантования. Для HardTanh удобно выбрать. Тогда граничные значения соответствуют максимальным квантованным числам.

еперь предположим, что наш вход — не вещественное число, а целочисленный результат матричной операции. У него есть собственный шаг квантования и при этом вещественный эквивалент. А это значит, что операции деквантования (перевода квантованных значений обратно в вещественные), активации и повторного квантования можно объединить в одну:

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

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

Казалось бы – хорошее решение. Но возникают две новые проблемы, которые мы рассмотрим далее.

Новые проблемы: ветвления и градиенты

Начнём с ветвлений.

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

В качестве понятного примера возьмём Res-блок (структурную единицу сетей ResNet) — хорошо известную и широко распространённую конструкцию с residual-связью. В обычном Res-блоке есть две ветви. В основной ветви сигнал проходит через два сверточных слоя (со сверткой 3x3). В residual-ветви он передаётся напрямую или через дополнительную свёртку (1x1), если нужно изменить размерность карты признаков. В конце ветви складываются.

Для вещественной сети это привычная операция. Для 4.6-битного исполнения всё сложнее. У двух ветвей могут быть разные масштабы (шаги квантования), а значит, мы не можем просто так сложить две ветви. Кроме того, в классическом Res-блоке в residual-ветви нет активации, как её нет и после второй свёртки основной ветви. Активация происходит после сложения. А значит мы снова сталкиваемся с проблемой промежуточных 32-битных карт признаков.

Мы предложили три варианта квантованного исполнения Res-блока. 

• В варианте A первая свёртка основной ветви работает как обычный 4.6-битный слой с HardTanh, но вторая свёртка и свёртка в residual-ветви возвращают результат в 32-битном формате; уже после сложения ветвей применяется активация и выход блока снова квантуется. 

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

• В варианте C вторая свёртка основной ветви и свёртка residual-ветви также квантуются, но так, чтобы минимизировать ошибку квантования и не сильно “порезать” хвосты распределения. При слиянии ветвей операции деквантования, сложения, активации и повторного квантования объединены в одну операцию без нормализации.

Другая проблема касается обучения сетей с HardTanh. Рассмотрим слой на этапе обучения, ещё до квантования: сначала выполняется линейное преобразование или свёртка, затем пакетная нормализация (batch norm), затем применяется HardTanh-активация. На вход активации приходит некоторое распределение значений, и от его масштаба зависит, в каком режиме будет работать слой.

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

Если распределение слишком узкое, почти все значения остаются в центральной линейной части HardTanh. Тогда активация почти ничего не меняет, а сеть теряет важную часть своей нелинейности. В пределе такая сеть начинает вести себя почти как линейная модель, что ограничивает её обобщающую способность.

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

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

Снова возвращаемся к теории

Задача выбора масштаба при инициализации нейронной сети не нова. Классические методы инициализации параметров в нейронных сетях Xavier [1] и Kaiming [2] как раз разрабатывались на основе анализа дисперсии сигнала и градиента при прохождении через сеть и обратном распространении ошибки.

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

Чтобы понять, что происходит с обучением, мы построили простую математическую модель слоя: матричное умножение, затем пакетная нормализация, затем HardTanh-активация. Это не полное описание всей сети, а минимальная схема, на которой удобно посмотреть, как меняются масштабы сигнала и градиента. Мы предполагаем, что значения перед HardTanh распределены нормально, их среднее равно нулю, а дисперсия задаётся параметром нормализации. Такое предположение оправдано в начале обучения, когда входы и веса независимы и работает центральная предельная теорема. После HardTanh распределение меняется: центральная часть проходит почти линейно, а хвосты обрезаются. Поэтому дисперсия выхода становится функцией двух параметров: ширины диапазона HardTanh и дисперсии после пакетной нормализации. Если эти параметры одинаково заданы во всех слоях, то дисперсия выхода после активации тоже остаётся одинаковой.

Аналогично смотрим на обратное распространение. Через линейную часть HardTanh градиент проходит, а через насыщенные участки — нет. Можно было бы ожидать, что основная опасность здесь в затухании градиента. Но расчёт показывает более интересную картину. Для слоя с квадратной матрицей весов и при дисперсия градиента на входе слоя оказывается примерно в 1.32 раза больше дисперсии градиента на выходе. Один слой сам по себе ещё не проблема. Но в глубокой сети такой множитель начинает накапливаться, и градиент может расти от слоя к слою при обратном распространении.

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

Экспериментируем с обучением

К этому моменту у нас уже есть два вопроса, которые нужно проверить экспериментально.

Первый — какой вариант квантованного Res-блока выбрать. Вариант A ближе к исходной архитектуре, но оставляет часть промежуточных данных в 32-битном формате. Вариант B одинаково обрабатывает все свёрточные слои в обеих ветвях, но требует более сложной операции объединения. Вариант C выглядит как компромисс: он позволяет отказаться от 32-битных промежуточных карт признаков на выходах ветвей, но остаётся ближе к обычному Res-блоку.

Второй вопрос — как инициализировать сеть. Из теоретического анализа мы получили две понятные стратегии для сравнения: обычную инициализацию, сохраняющую среднюю дисперсию сигнала, и инициализацию с адаптивными активациями, сохраняющую дисперсию градиента. 

Остаётся ещё один важный вопрос: как именно переводить сеть в 4.6-битный формат.

В задачах квантования обычно различают два больших подхода: PTQ и QAT. PTQ, или post-training quantization, применяется к уже обученной сети: мы берём готовые веса, подбираем параметры квантования по небольшой калибровочной выборке и получаем низкоразрядную модель без полноценного дообучения. Это удобно, когда нет доступа к обучающей выборке или нет ресурсов на повторное обучение. Но для низкоразрядных сетей такой путь часто оказывается слишком грубым: ошибка квантования велика, а у модели почти нет возможности к ней адаптироваться.

QAT, или quantization-aware training, устроен иначе. Сеть обучается с учётом будущего квантования: слои постепенно переводятся в низкоразрядный режим, а параметры донастраиваются так, чтобы компенсировать появляющуюся ошибку. Но QAT тоже можно организовать по-разному.

В экспериментах мы сравнивали несколько вариантов:

• PTQ: квантование по калибровочной выборке без полноценного дообучения.

• QAT-A: слой за слоем. Квантуем первый слой, дообучаем сеть, квантуем следующий, и так далее. При этом последний и предпоследний из квантованных на данный момент слоев учатся с помощью аппроксимации градиента. А все те что выше (ближе ко входу) – замораживаются. Шаг квантования для весов подбирается так, чтобы минимизировать их собственную ошибку квантования.

• QAT-B: нейрон за нейроном. То же, что и A, но вместо единовременного квантования целого слоя квантуется только несколько фильтров за раз (чтобы квантование шло плавнее).

• QAT-C: QAT с AdaQuant-подобной калибровкой. То же, что и A, но в момент квантования слоя веса настраиваются так, чтобы минимизировать ошибку квантования всего слоя (а не только весов). Это делается с помощью небольшой калибровочной выборки.

Такой набор вариантов позволяет проверить сразу несколько вещей: насколько PTQ уступает QAT, помогает ли постепенное квантование во время обучения и даёт ли калибровка перед квантованием слоя выигрыш по точности.

Эксперименты проводились для сети ResNet20 на датасете CIFAR-10. Результаты показаны в таблице.

Из таблицы видно несколько вещей.

Во-первых, эксперимент помог выбрать вариант блока. Лучший результат по точности показывает ResBlockA. ResBlockC отстаёт совсем немного, а ResBlockB в этих экспериментах оказался слабее двух других вариантов. Если учитывать расход оперативной памяти, для практики интереснее выглядит ResBlockC: в нём нет промежуточных 32-битных карт признаков на выходах ветвей.

Во-вторых, инициализация, лучше сохраняющая дисперсию градиента, для ResBlock-A и ResBlock-C оказывается лучше обычной инициализации, сохраняющей дисперсию сигнала.

В-третьих, PTQ заметно уступает вариантам с дообучением, поэтому в он не выглядит подходящим методом обучения 4.6-битных сетей.

Наконец, среди вариантов QAT лучше всего проявляет себя подход с калибровкой слоя перед его квантованием. Он даёт лучшие результаты для ResBlockA и ResBlockC при инициализации с адаптивными активациями.

Что мы узнали?

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

Выбор HardTanh в качестве активации помог закрыть часть этих вопросов. Он ограничивает диапазон значений и поэтому хорошо совмещается с квантованием: выход слоя можно сразу приводить к 4.6-битному формату, не сохраняя большой промежуточный тензор в 32-битном виде.

Но это решение потребовало дополнительной работы. Для Res-блоков пришлось выбрать подходящий вариант квантованного исполнения и реализовать специальную операцию объединения квантованных ветвей. Для обучения пришлось отдельно разобраться с дисперсиями сигнала и градиента, а также выбрать наиболее подходящий алгоритм квантования. Для HardTanh-сетей полезнее оказалось сохранять масштаб градиента, а для 4.6-битного квантования — использовать QAT с послойным переходом, калибровкой слоя перед квантованием и последующей заморозкой уже квантованных слоёв.

Где почитать научно про это?

Пост подготовлен по материалам наших научных статей:

• A. Trusov, E. Limonova, D. Nikolaev and V. V. Arlazarov, “4.6-bit Quantization for Fast and Accurate Neural Network Inference on CPUs,” Mathematics, vol. 12, no 5, pp. 651-1-651-16, 2024, https://doi.org/10.3390/math12050651

• A. V. Trusov, “Training 4.6-Bit Convolutional Neural Networks with a HardTanh Activation Function,” Pattern Recognit. Image Anal., vol. 35, no 1, pp. 44-64, 2025, https://doi.org/10.1134/S105466182470144X.