Новейшая методика распределенного обучения больших моделей глубокого обучения

В программной инженерии уменьшение времени цикла оказывает сверхлинейное влияние на прогресс. В современном глубоком обучении время цикла часто составляет часы или дни. Самый простой способ ускорить обучение, параллелизм данных, - распределить копии модели между графическими процессорами и машинами и заставить каждую копию вычислять потери на сегменте обучающих данных. Затем градиенты этих потерь можно накапливать, используя сервер с одним параметром или что-то более интересное, например, ring all-reduce (по умолчанию в pytorch). После обратного распространения градиентов повторите с обновленной моделью. Но по мере увеличения размера пакета градиенты могут стать нестабильными, что приведет к расхождению в обучении.

Недавняя статья сократила время обучения одной из самых больших и популярных языковых моделей, BERT, с 3 дней до 76 минут за счет более тонкого вычисления градиента, что позволило авторам масштабировать размер пакета от ~ 100X до 65K. . Хотя в документе описывается его вклад в отношении его предыдущей итерации, LARS, в этом посте я объясню необходимые предпосылки, начиная со стохастического градиентного спуска (SGD), на котором построен LARS.

Во-первых, обзор основ. Напомним, что для SGD стандартным правилом обновления является

Разбивая это на части, говорится, что веса новой модели - это существующие, минус скорость обучения, умноженная на градиент существующего веса на примере или среднее значение по партии. Проблемы:

  1. Мы не знаем истинного градиента, у нас есть только зашумленные оценки. Шум может вызывать ложные обновления
  2. Шкала градиента не нормализована, поэтому наше обновление может исчезнуть или взорваться со временем.

Адам

В Adam мы храним скользящее среднее градиентов и их дисперсию:

где 𝓂 - скользящее среднее, 𝓋 - скользящая нецентрированная дисперсия, β₁ - постоянная интерполяции для среднего, β₂ - постоянная интерполяции для нецентрированной дисперсии, а ∇L - градиент потерь. Скобки в показателях степени означают, что это не показатель степени, а временной шаг. Это выглядит пугающе, но важно отметить, что 𝓂 и 𝓋 - это просто линейные интерполяции (β * x ₀ + (1 - β) * x ₁ ) градиентов и их дисперсий, что дает нам скользящее среднее для каждого. Чем выше бета, тем меньше мы обновляем скользящее среднее для каждой новой выборки, тем самым сглаживая нашу оценку среднего и дисперсии градиента между партиями. Вот визуализация того, насколько сглаживается шумный набор данных для разных бета-версий.

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

Проблема со скользящими средними, указанными выше, заключается в том, что при первой инициализации алгоритма скользящие средние равны 0. Это приводит к тому, что сводная статистика становится ближе к 0, чем должна быть для первых двух временных шагов, если бета близка к 1, потому что мы взяв большую часть массы из предыдущего шага. Этот эффект особенно заметен на графике бета = 0,99 выше.

Мы исправляем это путем удаления

Что приводит к:

Проблема никуда не делась, но стало намного лучше.

Чтобы подставить некоторые числа, если β = 0,9, на первой итерации debias умножит значение на 1 / 1–0,9¹ = 10. Затем, когда мы линейно интерполируем, β𝓍₀ + (1 - β) 𝓍₁, первый член β𝓍₀ = 0. Коэффициент дебиаса, 10, компенсирует (1 - β) = 0,1 во втором члене, поэтому мы полностью используем новое значение 𝓍₁. Всего через несколько шагов коэффициент дебиаса сходится к 1. Графики ниже показывают, сколько шагов нужно, чтобы термин дебиас исчез (обратите внимание на разницу по оси y):

Теперь последнее обновление параметров

В числителе указано «для каждого параметра сделайте шаг в направлении градиента для этого параметра». Знаменатель говорит: «нормализовать шаг по его стандартному отклонению».

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

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

Это отношение m / sqrt (v) может выглядеть как μ / σ, что является отношением сигнал / шум, но эта интерпретация применима только к скалярам.

ЛАРС

По мере увеличения размера пакета количество итераций за эпоху уменьшается. Чтобы сойтись за одно и то же количество итераций набора данных, мы можем компенсировать это, увеличив скорость обучения. Однако по мере увеличения скорости обучения обучение становится все более нестабильным. SOTA должна была использовать разминку скорости обучения, но это помогает только до определенного момента, когда обучение все равно начнет расходиться. Разминка была исправлением реальной проблемы: градиенты должны быть шумными.

Авторы Layerwise Adaptive Rate Scaling (LARS) объясняют свой прием для решения этой проблемы:

Для анализа устойчивости обучения с большими LR мы измерили соотношение между нормой весов слоев и нормой обновления градиентов. Мы заметили, что если это соотношение слишком велико, тренировка может стать нестабильной. С другой стороны, если соотношение слишком мало, веса меняются недостаточно быстро.

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

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

Числитель - это норма весов, потому что по мере того, как сети углубляются, становится важным иметь нулевое среднее значение, веса единичной дисперсии (ZMUV). Это связано с тем, что на каждом слое эти веса умножаются вместе, поэтому, если он отличается от ZMUV, значения могут взорваться или исчезнуть. Когда вес небольшой, мы делаем маленький шаг. Когда вес большой, мы делаем больший шаг. В сочетании со снижением веса это помогает нам стабильно приближаться к весам ZMUV.

Давайте разберемся, что здесь происходит. В начале обучения слои должны выводить ZMUV, поэтому числитель выше будет равен 0 или близок к нему. Любые шаги, которые мы предпримем, скорее всего, будут небольшими. Напротив, знаменатель, вероятно, будет большим, поскольку, когда все не так, градиенты большие. Таким образом мы естественным образом разминаемся по мере увеличения веса. Когда мы приближаемся к 0 потерям, градиенты будут небольшими, поэтому коэффициент доверия будет поддерживать скорость обучения до 10 раз (из-за ограничения) выше, чем без коэффициента доверия, не позволяя нам отказаться от слишком раннего достижения оптимума.

ЯГНЕНОК

LAMB означает «Послойный оптимизатор адаптивных моментов для пакетного обучения». Вносит несколько небольших изменений в LARS.

  1. Если числитель (r₁ ниже) или знаменатель (r₂ ниже) коэффициента доверия равен 0, используйте вместо этого 1. Этот раздел было трудно читать, поэтому я основывал это на некотором коде.
  2. Фиксация снижения веса: в LARS знаменатель коэффициента доверия | ∇L | + β | w |, тогда как в LAMB это | ∇L + β w |. Это сохраняет больше информации.
  3. Вместо использования правила обновления SGD они используют правило обновления Адама.
  4. Уменьшите коэффициент доверия до 10.

Таким образом, в целом коэффициент доверия в LAMB составляет

Последняя строка - это послойное правило обновления LAMB. 𝓇₂ - норма правила обновления Адама с уменьшением веса, ηᴸ - скорость обучения по уровням, скорректированная с учетом коэффициента доверия. В целом этот метод можно резюмировать как LARS, примененный к Адаму, поскольку он просто умножает старый шаг обновления на коэффициент доверия.

Авторы не сообщают, улучшает ли LAMB эффективность обучения ImageNet по сравнению с LARS, и они не сравнивают LARS с LAMB для BERT, поэтому немного сложно сказать, насколько эти изменения имеют значение, но реализовать довольно просто.

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

Чтобы лучше понять, что происходит, я реализовал LAMB в Pytorch. Я провел кучу экспериментов с MNIST и обнаружил, что там, где Адам расходится, LAMB продолжает пыхтеть. Я выбрал MNIST, потому что он достаточно мал, чтобы его можно было попробовать на ЦП, но это означает, что мы не видим никаких улучшений сходимости. Я скоро опубликую еще один пост, посвященный LAMB в применении к большим трансформерам!

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

Ниже я сравниваю Adam (синий внизу) и LAMB (красный внизу) со скоростью обучения 0,01 и бета 0,9, 0,99. Они довольно похожи, но LAMB лучше обобщает точность тестов.

Чтобы выяснить, что происходит под капотом, я хотел посмотреть на послойные компоненты отношения доверия, поэтому я регистрировал каждое значение r, r₁ и r₂ в виде гистограммы после каждых нескольких пакетов. Для Адама я вычисляю значения, но нигде их не использую.

Как интерпретировать диаграмму ниже: ось Y показывает, какой временной шаг, первый вверху, ось X - сегменты гистограммы, ось Z - частота гистограммы.

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

Для следующего эксперимента я сравнил LAMB с самим собой при скорости обучения 0,1 и 0,01. Адам обычно сходится при скорости обучения 0,01, а при 0,1 не учится вообще, поэтому я не буду сравнивать это здесь. Слева (синий) скорость обучения = 0,01, справа (зеленый) скорость обучения = 0,1. Справа он почти мгновенно сходится во время разминки, но затем веса нескольких слоев начинают взрываться (см. Разницу в шкале оси X), и они расходятся.

Чтобы справиться с бегством веса, я добавил снижение веса на 0,01 внизу справа. Тренировки не расходились! Как правило, коэффициент доверия продолжал медленно учиться и составлял менее 1, тогда как в более комфортном режиме, показанном выше слева, он достигал 4,5.

Резюме

Ванильный SGD становится нестабильным по мере увеличения скорости обучения.

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

Адам модулирует обновления с помощью отклоненных средств, нормализованных по отклоненным отклонениям.

LAMB регулирует скорость обучения Адама более точным послойным, обрезанным коэффициентом доверия.

Сочетание всех этих методов позволяет нам тренироваться на больших партиях с высокой скоростью обучения, сокращая время перерыва в 100 раз для BERT!

Спасибо Ярославу Булатову и Саре Джейн Хонг за правки и Джереми Ховарду /fast.ai, часть 2, за вдохновение