2025-11-13T04:07:09.837900

Optimal Quantization for Matrix Multiplication

Ordentlich, Polyanskiy
Recent work in machine learning community proposed multiple methods for performing lossy compression (quantization) of large matrices. This quantization is important for accelerating matrix multiplication (main component of large language models), which is often bottlenecked by the speed of loading these matrices from memory. Unlike classical vector quantization and rate-distortion theory, the goal of these new compression algorithms is to be able to approximate not the matrices themselves, but their matrix product. Specifically, given a pair of real matrices $A,B$ an encoder (compressor) is applied to each of them independently producing descriptions with $R$ bits per entry. These representations subsequently are used by the decoder to estimate matrix product $A^\top B$. In this work, we provide a non-asymptotic lower bound on the mean squared error of this approximation (as a function of rate $R$) for the case of matrices $A,B$ with iid Gaussian entries. Algorithmically, we construct a universal quantizer based on nested lattices with an explicit guarantee of approximation error for any (non-random) pair of matrices $A$, $B$ in terms of only Frobenius norms $\|\bar{A}\|_F, \|\bar{B}\|_F$ and $\|\bar{A}^\top \bar{B}\|_F$, where $\bar{A},\bar{B}$ are versions of $A,B$ with zero-centered columns, respectively. For iid Gaussian matrices our quantizer achieves the lower bound and is, thus, asymptotically optimal. A practical low-complexity version of our quantizer achieves performance quite close to optimal. In addition, we derive rate-distortion function for matrix multiplication of iid Gaussian matrices, which exhibits an interesting phase-transition at $R\approx 0.906$ bit/entry, showing necessity of Johnson-Lindestrauss dimensionality reduction (sketching) in the low-rate regime.
academic

Оптимальное квантование для матричного умножения

Основная информация

  • ID статьи: 2410.13780
  • Название: Optimal Quantization for Matrix Multiplication
  • Авторы: Or Ordentlich (Hebrew University of Jerusalem), Yury Polyanskiy (MIT)
  • Классификация: cs.IT cs.AI cs.CL cs.LG math.IT
  • Дата публикации: Октябрь 2024 г. (препринт arXiv)
  • Ссылка на статью: https://arxiv.org/abs/2410.13780

Аннотация

В данной работе проводится глубокое исследование проблемы квантования при крупномасштабном матричном умножении. В отличие от традиционного векторного квантования, целью исследования является не приближение самих матриц, а приближение их матричного произведения. Для двух вещественных матриц A и B кодировщик независимо сжимает каждую матрицу, описывая каждый элемент R битами, затем декодировщик использует эти сжатые представления для оценки матричного произведения A⊤B. Статья предоставляет неасимптотические нижние границы среднеквадратической ошибки для матриц с независимыми одинаково распределёнными гауссовыми элементами, конструирует универсальный квантователь на основе вложенных решёток и обнаруживает интересный фазовый переход при R ≈ 0,906 бит/элемент, указывающий на необходимость техники понижения размерности Johnson-Lindenstrauss при низких скоростях кодирования.

Исследовательский контекст и мотивация

Определение проблемы

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

Практические требования

Для больших языковых моделей авторы оценили требуемые скорости квантования:

  • На этапе генерации процессору требуется скорость квантования 1-3 бит/элемент для полного использования вычислительных ресурсов
  • На этапе предзаполнения для небольших LLM, работающих на быстрых GPU, требуется примерно 11,7 бит/элемент

Ограничения существующих методов

  1. Классическое векторное квантование: независимое квантование матриц A и B с последующим вычислением произведения квантованных матриц приводит к ошибке O(n²)
  2. Методы скетчирования: хотя и обеспечивают несмещённую оценку, дисперсия остаётся O(n²)
  3. Детерминированные квантователи: для векторов на сфере существует нижняя граница Ω(n²)

Основные вклады

  1. Теоретические нижние границы: предоставляются неасимптотические нижние границы квантования матричного умножения для матриц с независимыми одинаково распределёнными гауссовыми элементами
  2. Универсальный квантователь: конструируется универсальный квантователь на основе вложенных решёток с явными гарантиями ошибки для произвольных матриц
  3. Асимптотическая оптимальность: доказывается, что предложенный квантователь достигает нижней границы для матриц с независимыми одинаково распределёнными гауссовыми элементами, следовательно, является асимптотически оптимальным
  4. Фазовый переход: обнаруживается фазовый переход при R ≈ 0,906 бит/элемент, раскрывающий необходимость понижения размерности при низких скоростях кодирования
  5. Практические алгоритмы: предоставляются реализации с низкой сложностью, обеспечивающие производительность, близкую к оптимальной

Подробное описание методов

Определение задачи

Для матриц A ∈ R^(n×a) и B ∈ R^(n×b) целью является разработка кодировщиков f₁: R^(n×a) → 2^(naR) и f₂: R^(n×b) → 2^(nbR) и декодировщика g таких, что:

E∥A⊤B−g(f1(A),f2(B))∥F2E\|A^⊤B - g(f_1(A), f_2(B))\|_F^2

минимизируется, где каждый элемент матрицы описывается R битами.

Ключевая функция Γ(R)

Статья определяет критическую функцию скорость-искажение:

Γ(R)={1−(1−(2⋅2−2R∗−2−4R∗))RR∗R<R∗2⋅2−2R−2−4RR≥R∗\Gamma(R) = \begin{cases} 1 - \left(1 - (2 \cdot 2^{-2R^*} - 2^{-4R^*})\right) \frac{R}{R^*} & R < R^* \\ 2 \cdot 2^{-2R} - 2^{-4R} & R \geq R^* \end{cases}

где R* ≈ 0,906 является решением уравнения неподвижной точки R = ½log₂(1 + 4R ln 2).

Схема квантования на основе вложенных решёток

1. Квантование внутреннего произведения (базовый строительный блок)

Для ρ-коррелированных векторов U, V на сфере используются вложенные решётки Λc ⊂ Λf:

Процесс кодирования:

  • Добавляются независимые векторы дрожания Z₁, Z₂ к U и V соответственно
  • Квантование к точкам тонкой решётки Λf
  • Вывод представления смежного класса в грубой решётке Λc

Процесс декодирования:

  • Восстановление квантованной точки из смежного класса
  • Удаление дрожания
  • Вычисление оценки внутреннего произведения

2. Квантование матричного умножения

Этап предварительной обработки:

  1. Центрирование нулём: вычисляются Ā = A - (1/n)1·1^⊤A, B̄ = B - (1/n)1·1^⊤B
  2. Квантование норм: высокоточное описание средних значений и норм каждого столбца
  3. Случайное вращение: применение одной и той же ортогональной матрицы S к Ā и B̄

Этап квантования:

  • Применение квантователя внутреннего произведения к каждому столбцу после вращения
  • Использование параметров временного разделения κ и параметра масштабирования MMSE α

Технические инновации

  1. Техника дрожания: делает ошибку квантования независимой от входных данных, избегая ошибки O(n²) детерминированных квантователей
  2. Структура вложенных решёток: обеспечивает хорошую производительность квантования при конечной скорости кодирования
  3. Временное разделение: достижение оптимальной производительности при низких скоростях кодирования через понижение размерности
  4. Случайное вращение: преобразование произвольных векторов в равномерное распределение на сфере, облегчающее анализ

Экспериментальная установка

Эксперименты по проверке теории

Генерация данных:

  • Матрицы A, B ∈ R^(n×n) с независимыми одинаково распределёнными элементами N(0,1)
  • n = 3 × 2¹¹

Детали реализации:

  • Базовая решётка: решётка D₃ (трёхмерная)
  • Коэффициент вложения: q = 6
  • Размер таблицы поиска: < 64 КБ (может поместиться в кэш L1)
  • Эффективная скорость: ≈ 3,015 бит/символ

Методы сравнения

  • Трёхбитный скалярный квантователь (нормализация ℓ∞)
  • Теоретическая нижняя граница Γ(R)

Результаты экспериментов

Основные результаты

Сравнение производительности:

  • Предложенный метод: 1/n³ ∥Â⊤B - A⊤B∥²F ≈ 0,0593
  • Трёхбитное скалярное квантование: ≈ 0,1668 (примерно 3-кратная разница)
  • Теоретическая нижняя граница: Γ(3,015) = 0,0304

Ключевые находки:

  1. Схема на основе решётки D₃ значительно превосходит скалярное квантование
  2. Производительность близка к теоретически оптимальной (примерно 2-кратная разница)
  3. По сравнению с ростом n разница в производительности будет дополнительно сокращаться

Анализ сложности

Сложность кодирования: O(n log n) (с использованием быстрого преобразования Адамара) Сложность декодирования: O(1) (с использованием таблицы поиска) Накладные расходы на хранение: каждый квантователь требует O(log n) дополнительных битов для описания масштабирующих коэффициентов

Связанные работы

Случайная линейная алгебра

  • Матричное умножение методом Монте-Карло (MCMM): случайная выборка строк для приближения
  • Локально чувствительное хеширование (LSH): низкомерные скетчи для косинусного сходства
  • Ограничения: относительная ошибка растёт с ∥A∥²F∥B∥²F/∥A⊤B∥²F

Квантование нейронных сетей

  • Квантование после обучения: методы OPTQ, GPTQ и др.
  • Техники вращения: QuIP, QuaRot используют преобразование Адамара
  • Квантование решётками: QUIP# использует решётку E₈ для квантования весов

Теория информации

  • Распределённое сжатие: сжатие для вычисления линейных функций
  • Проектирование кодовых книг: коды Вороного и коды вложенных решёток

Выводы и обсуждение

Основные выводы

  1. Оптимальность: для матриц с независимыми одинаково распределёнными гауссовыми элементами предложенная схема достигает информационно-теоретической нижней границы
  2. Универсальность: обеспечивает явные гарантии производительности для произвольных матриц
  3. Фазовый переход: R* ≈ 0,906 бит/элемент является критическим порогом
  4. Практичность: реализация с низкой сложностью обеспечивает производительность, близкую к теоретической

Ограничения

  1. Общая случайность: требуется, чтобы кодировщик и декодировщик делили общее случайное зерно
  2. Условия на матрицы: требуется ограниченность элементов матрицы (M = n^(10^22000))
  3. Высокомерные решётки: теоретическая оптимальность требует высокомерных "хороших" решёток, на практике используются произведения низкомерных решёток
  4. Детерминированные схемы: неясно, существуют ли оптимальные детерминированные схемы без необходимости общей случайности

Направления будущих исследований

  1. Произведения нескольких матриц: расширение на произведения k>2 матриц
  2. Другие метрики расстояния: метрики, отличные от MSE, такие как расхождение Кульбака-Лейблера
  3. Детерминированные квантователи: исследование схем без необходимости общей случайности
  4. Применение в глубоких сетях: развёртывание и оптимизация в реальных LLM

Глубокая оценка

Преимущества

  1. Теоретическая строгость: предоставляется полный информационно-теоретический анализ, включая верхние и нижние границы
  2. Практическая ценность: решает реальное узкое место в выводе LLM
  3. Технические инновации: умелое сочетание квантования решётками, случайного вращения и временного разделения
  4. Универсальность: не зависит от предположений о конкретном распределении матриц

Недостатки

  1. Сложность: теоретический анализ довольно сложен, практическая реализация требует нескольких технических компонентов
  2. Константные множители: хотя асимптотически оптимально, константы при конечных выборках могут быть значительными
  3. Адаптация к оборудованию: требуется оптимизация реализации для различных аппаратных платформ
  4. Масштабируемость: расширение с двух матриц на произведения нескольких матриц нетривиально

Влияние

Теоретический вклад:

  • Установление информационно-теоретических основ квантования матричного умножения
  • Раскрытие фазового перехода и необходимости понижения размерности
  • Предоставление важного теоретического ориентира для данной области

Практическое применение:

  • Предоставление новых теоретических указаний для квантования LLM
  • Последующая работа NestQuant достигла SOTA производительности на реальных LLM
  • Предоставление теоретического обоснования для проектирования аппаратных ускорителей

Применимые сценарии

  1. Вывод больших языковых моделей: совместное квантование весов и активаций
  2. Граничные вычисления: матричные операции в среде с ограниченной памятью
  3. Распределённые вычисления: матричное умножение с ограничениями на коммуникацию
  4. Научные вычисления: крупномасштабные задачи численной линейной алгебры

Библиография

Статья цитирует 44 связанные работы, охватывающие теорию информации, теорию решёток, случайную линейную алгебру и квантование нейронных сетей. Особого внимания заслуживают:

  • Монография Zamira по кодированию решётками, предоставляющая теоретическую основу
  • Пионерская работа Erez и Zamir по вложенным решёткам
  • Недавние методы квантования LLM, такие как OPTQ, QuIP и др.
  • Результаты теории случайных матриц и сферической геометрии

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