Sparse Universal Transformer. Shawn Tan, Yikang Shen, Zhenfang...
Sparse Universal Transformer
Shawn Tan, Yikang Shen, Zhenfang Chen, Aaron Courville, Chuang Gan
Статья: https://arxiv.org/abs/2310.07096
Развитие моего любимого Universal Transformer, UT (https://t.me/gonzo_ML/62, https://t.me/gonzo_ML/90, https://moocaholic.medium.com/adaptive-computation-time-act-in-neural-networks-3-3-99452b2eff18) и более глобально темы про Adaptive Computation Time (ACT). Теперь есть Sparse Universal Transformer, SUT -- это UT + Sparse Mixture of Experts (SMoE), про MoE тоже было много (https://t.me/gonzo_ML/472).
Напомню в двух словах, в чём идея UT. Bместо множества разных слоёв с разными весами (как было в обычном трансформере, Vanilla Transformer, VT), в UT есть один слой, который можно применять многократно -- его выход снова отправляется ему же на вход и таким образом входные эмбеддинги итеративно обновляются. Для контроля этого процесса есть специальный отдельный предсказатель “готовности” выхода. В UT это делалось через предсказание pondering value, и процесс останавливался, когда кумулятивно это значение было больше заданного порога. Работало это на уровне каждого отдельного токена, так что некоторые (“сложные”) токены могли обрабатываться долго, а другие (“простые”) -- быстро.
После выхода работы про UT в 2018-м накопилось сколько-то результатов про его эффективность. В известной работе про иерархии Хомского (https://t.me/gonzo_ML/1049) было показано, что у VT есть проблемы генерализации, которых нет у некоторых других типов архитектур. Другие работы типа “The Devil is in the Detail: Simple Tricks Improve Systematic Generalization of Transformers” (https://arxiv.org/abs/2108.12284) показали, что UT генерализует лучше. Но со скейлингом у UT немного посложнее, и одна работа (https://arxiv.org/abs/2104.06022) показала, что UT нужно больше обучения и памяти на WMT, а другая (https://arxiv.org/abs/2207.10551) получила кривые скейлинга, где это видно.
Вышеупомянутая архитектура Mixture-of-Experts (MoE) может быть использована для уменьшения вычислительной сложности больших моделей. В ней в каждый момент активируется только часть сети, например, несколько “экспертов” (блоков сети) из множества. Когда из всех возможных экспертов (E) выбираются не все (k), это называют Sparse Mixture of Experts (SMoE). Широко известные классические примеры SMoE это Switch Transformer и GShard (https://t.me/gonzo_ML/472).
В текущей работе про SUT применяют SMoE отдельно к Multi-Head Self-Attention (MHSA) и к Feed-Forward (FF) слоям. В FF как обычно в SMoE вместо одного слоя появляется набор из множества FF-слоёв (экспертов), из которых выбирается top-k экспертов. Для MHSA используется Mixture of Multihead Attention (MoMHA) из работы “Mixture of Attention Heads: Selecting Attention Heads Per Token” (https://arxiv.org/abs/2210.05144). Здесь выбирается top-k голов внимания из всех доступных. Выбором экспертов занимается специальная gating network, выдающая распределение над экспертами, из этого распределения и выбирается топ экспертов. В итоге оба этих SMoE образуют блок SUT, и этот блок переиспользуется для всех слоёв нового трансформера.
Для обучения всей этой машинерии нужны дополнительные лоссы. Здесь при обучении максимизируется специальная unsupervised версия Mutual Information Maximization лосс для обучения gating network. Этот лосс позволяет с одной стороны получать острое распределение для выбора экспертов, а с другой балансировать загрузку экспертов, чтобы все выбирались равномерно на горизонте батча.
Shawn Tan, Yikang Shen, Zhenfang Chen, Aaron Courville, Chuang Gan
Статья: https://arxiv.org/abs/2310.07096
Развитие моего любимого Universal Transformer, UT (https://t.me/gonzo_ML/62, https://t.me/gonzo_ML/90, https://moocaholic.medium.com/adaptive-computation-time-act-in-neural-networks-3-3-99452b2eff18) и более глобально темы про Adaptive Computation Time (ACT). Теперь есть Sparse Universal Transformer, SUT -- это UT + Sparse Mixture of Experts (SMoE), про MoE тоже было много (https://t.me/gonzo_ML/472).
Напомню в двух словах, в чём идея UT. Bместо множества разных слоёв с разными весами (как было в обычном трансформере, Vanilla Transformer, VT), в UT есть один слой, который можно применять многократно -- его выход снова отправляется ему же на вход и таким образом входные эмбеддинги итеративно обновляются. Для контроля этого процесса есть специальный отдельный предсказатель “готовности” выхода. В UT это делалось через предсказание pondering value, и процесс останавливался, когда кумулятивно это значение было больше заданного порога. Работало это на уровне каждого отдельного токена, так что некоторые (“сложные”) токены могли обрабатываться долго, а другие (“простые”) -- быстро.
После выхода работы про UT в 2018-м накопилось сколько-то результатов про его эффективность. В известной работе про иерархии Хомского (https://t.me/gonzo_ML/1049) было показано, что у VT есть проблемы генерализации, которых нет у некоторых других типов архитектур. Другие работы типа “The Devil is in the Detail: Simple Tricks Improve Systematic Generalization of Transformers” (https://arxiv.org/abs/2108.12284) показали, что UT генерализует лучше. Но со скейлингом у UT немного посложнее, и одна работа (https://arxiv.org/abs/2104.06022) показала, что UT нужно больше обучения и памяти на WMT, а другая (https://arxiv.org/abs/2207.10551) получила кривые скейлинга, где это видно.
Вышеупомянутая архитектура Mixture-of-Experts (MoE) может быть использована для уменьшения вычислительной сложности больших моделей. В ней в каждый момент активируется только часть сети, например, несколько “экспертов” (блоков сети) из множества. Когда из всех возможных экспертов (E) выбираются не все (k), это называют Sparse Mixture of Experts (SMoE). Широко известные классические примеры SMoE это Switch Transformer и GShard (https://t.me/gonzo_ML/472).
В текущей работе про SUT применяют SMoE отдельно к Multi-Head Self-Attention (MHSA) и к Feed-Forward (FF) слоям. В FF как обычно в SMoE вместо одного слоя появляется набор из множества FF-слоёв (экспертов), из которых выбирается top-k экспертов. Для MHSA используется Mixture of Multihead Attention (MoMHA) из работы “Mixture of Attention Heads: Selecting Attention Heads Per Token” (https://arxiv.org/abs/2210.05144). Здесь выбирается top-k голов внимания из всех доступных. Выбором экспертов занимается специальная gating network, выдающая распределение над экспертами, из этого распределения и выбирается топ экспертов. В итоге оба этих SMoE образуют блок SUT, и этот блок переиспользуется для всех слоёв нового трансформера.
Для обучения всей этой машинерии нужны дополнительные лоссы. Здесь при обучении максимизируется специальная unsupervised версия Mutual Information Maximization лосс для обучения gating network. Этот лосс позволяет с одной стороны получать острое распределение для выбора экспертов, а с другой балансировать загрузку экспертов, чтобы все выбирались равномерно на горизонте батча.
Источник: gonzo-обзоры ML статей
2023-10-22 22:02:53