NeuZip: Эффективная компрессия нейронных сетей для обучения и вывода

Современные достижения в области глубокого обучения во многом связаны с увеличением числа параметров нейронных сетей. Модели с десятками миллиардов параметров, такие как GPT-3, стали стандартом в области обработки естественного языка и других задач. Однако, несмотря на впечатляющий рост производительности, существует серьезная проблема: доступная память на устройствах, таких как графические процессоры (GPU), не успевает за увеличением размера моделей. Например, на момент 2024 года максимальный объем памяти на GPU составляет 80 ГБ, что значительно меньше, чем требуется для обучения современных моделей с сотнями миллиардов параметров.

В этом контексте необходимо разрабатывать методы, которые позволят уменьшить объем памяти, необходимый для обучения и вывода нейронных сетей, не ухудшая их производительность. В данной статье мы обсудим NeuZip — новый алгоритм компрессии весов, который позволяет достигать эффективного использования памяти без потери точности.

Проблема и подход

Ограничения по памяти

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

Однако для параметров модели до недавнего времени не существовало эффективных методов, позволяющих сохранить память при обучении. К примеру, метод QLoRA, который использует 4-битные данные для предварительно обученной модели, позволяет значительно сократить объем памяти, но при этом накладывает ограничения на изменения модели, что может снизить ее производительность.

NeuZip: Новый подход к компрессии

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

NeuZip использует асимметричную числовую систему (ANS) для компрессии экспонент, что позволяет достичь высокой пропускной способности на параллельных вычислительных устройствах, таких как GPU. При этом компрессия является без потерь, что означает, что качество модели сохраняется на прежнем уровне.

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

Технические детали NeuZip

Энтропия параметров нейронной сети

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

Без потерь: Компрессия экспонент

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

С потерями: Трuncирование мантиссы

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

Эксперименты и результаты

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

Обучение и вывод

В экспериментах с предобучением языковых моделей, таких как Llama-3 и GPT-Neo, NeuZip продемонстрировал возможность обучения моделей с 13 миллиардами параметров на устройствах с ограниченной памятью (менее 20 ГБ), что делает его особенно привлекательным для исследователей с ограниченными ресурсами.

Сравнение с другими методами

NeuZip был сравнен с другими методами компрессии, такими как QLoRA и LOMO. Результаты показали, что NeuZip обеспечивает лучшее соотношение между памятью и производительностью, находясь на так называемой "передней границе Парето" в области компромисса между памятью и производительностью.

Заключение

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

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