LayerShuffle: Enhancing Robustness in Vision Transformers by Randomizing Layer Execution Order
Matthias Freiberger, Peter Kun, Anders Sundnes Løvlie, Sebastian Risi
Статья: https://arxiv.org/abs/2407.04513
Недавно вышли две свежие прикольные работы про творческий подход к вычислению трансформерных слоёв. Это перекликается с нежно любимой мной темой про Adaptive Computation Time (ACT) и вообще динамическими вычислениями. Про трансформеры c ACT было тут (https://moocaholic.medium.com/adaptive-computation-time-act-in-neural-networks-3-3-99452b2eff18), но вообще это целая серия постов (в канале она начиналась тут https://t.me/gonzo_ML/71).
Сегодня первая интересная работа -- про выкидывание части трансформерных слоёв во время инференса. Если сеть будет устойчива к таким событиям, то это открывает интересные возможности для выполнения вычислений в нестабильной распределённой среде, когда часть вычислений может выпасть или произойти в ином порядке. Способность продолжать работу в таких условиях и не терять сильно в качестве при этом довольно интересна.
Это идейно похоже на Dropout, где выпадает часть нейронов в слое, но делается уровнем выше. В 2019-м одни из моих любимых авторов опубликовали вариант structured dropout под названием LayerDrop (https://arxiv.org/abs/1909.11556), позволявший делать прунинг части слоёв трансформера при инференсе, при этом после прунинга он приводил к моделям с качеством выше, чем если бы модели такого же размера обучались с нуля или даже были получены дистилляцией (!). В этом смысле интересно было бы посмотреть на вариант той же Gemma 2, где малую модель не дистиллировали, а дропали из большой.
LayerShuffle идёт дальше и не просто дропает слои при инференсе, а перемешивает их, меняя порядок вычисления. Делают это на Vision Transformer’ах (ViT, https://t.me/gonzo_ML/434), но, кажется, это неважно. Рассматриваются несколько вариантов процедуры.
В простом случае, LayerShuffle, во время обучения порядок слоёв случайно меняется для каждого батча. Таким образом слои учатся не затачиваться на выход конкретного другого слоя.
Во втором варианте, LayerShuffle-position, каждый слой также получает эмбеддинг текущей позиции этого слоя (эти дополнительные 32 чиселки конкатенируются с эмбеддингом, прилетающим на вход).
Третий вариант, LayerShuffle-predict, пытается предсказать свою текущую позицию по своему же выходному эмбеддингу. С кроссэнтропийным лоссом при этом, что немного странно, ибо сильный промах и слабый промах в позиции получается штрафуются одинаково.
Всё проверяют на ImageNet2012 с предобученным ViT-B/16 с HuggingFace, который дообучали по описанным процедурам. Стандартный трансформер при переключении в режим случайного порядка слоёв роняет качество с 82% почти в ноль, а варианты LayerShuffle падают примерно с 75% до 63% (интересно, какая была бы разница, если бы модель обучали с нуля по предложенному подходу, а не файнтюнили с классически обученной). LayerShuffle-position лидирует с небольшим отрывом от простого LayerShuffle.
Получается, стандартный трансформер катастрофически неустойчив к смене порядка вычисления слоёв, что было в целом ожидаемо. Интересно, что вариант LayerShuffle-position с передаваемой позицией не очень-то сильно что-то улучшает, все важные данные, видимо, есть и так.
Отдельно проверили эффект прунинга во время инференса. Качество работы примерно соответствует LayerDrop с p=0.2, хотя модель на выкидывание слоёв не обучалась. Если дополнительно к прунингу ещё и перемешивать, то LayerDrop тоже уходит в ноль, а LayerShuffle держится. Получается, если выбирать из этих двух, то последний даёт более устойчивые модели.
Сделали визуализацию выходов слоёв через UMAP, не передавая в алгоритм информацию о позиции слоя. Видно, что с одной стороны явного прерывистого разделения эмбеддингов нет, но таки они группируются по своей реальной позиции, особенно заметно для эмбеддингов близких к началу, но впрочем видно и для конечных.
Matthias Freiberger, Peter Kun, Anders Sundnes Løvlie, Sebastian Risi
Статья: https://arxiv.org/abs/2407.04513
Недавно вышли две свежие прикольные работы про творческий подход к вычислению трансформерных слоёв. Это перекликается с нежно любимой мной темой про Adaptive Computation Time (ACT) и вообще динамическими вычислениями. Про трансформеры c ACT было тут (https://moocaholic.medium.com/adaptive-computation-time-act-in-neural-networks-3-3-99452b2eff18), но вообще это целая серия постов (в канале она начиналась тут https://t.me/gonzo_ML/71).
Сегодня первая интересная работа -- про выкидывание части трансформерных слоёв во время инференса. Если сеть будет устойчива к таким событиям, то это открывает интересные возможности для выполнения вычислений в нестабильной распределённой среде, когда часть вычислений может выпасть или произойти в ином порядке. Способность продолжать работу в таких условиях и не терять сильно в качестве при этом довольно интересна.
Это идейно похоже на Dropout, где выпадает часть нейронов в слое, но делается уровнем выше. В 2019-м одни из моих любимых авторов опубликовали вариант structured dropout под названием LayerDrop (https://arxiv.org/abs/1909.11556), позволявший делать прунинг части слоёв трансформера при инференсе, при этом после прунинга он приводил к моделям с качеством выше, чем если бы модели такого же размера обучались с нуля или даже были получены дистилляцией (!). В этом смысле интересно было бы посмотреть на вариант той же Gemma 2, где малую модель не дистиллировали, а дропали из большой.
LayerShuffle идёт дальше и не просто дропает слои при инференсе, а перемешивает их, меняя порядок вычисления. Делают это на Vision Transformer’ах (ViT, https://t.me/gonzo_ML/434), но, кажется, это неважно. Рассматриваются несколько вариантов процедуры.
В простом случае, LayerShuffle, во время обучения порядок слоёв случайно меняется для каждого батча. Таким образом слои учатся не затачиваться на выход конкретного другого слоя.
Во втором варианте, LayerShuffle-position, каждый слой также получает эмбеддинг текущей позиции этого слоя (эти дополнительные 32 чиселки конкатенируются с эмбеддингом, прилетающим на вход).
Третий вариант, LayerShuffle-predict, пытается предсказать свою текущую позицию по своему же выходному эмбеддингу. С кроссэнтропийным лоссом при этом, что немного странно, ибо сильный промах и слабый промах в позиции получается штрафуются одинаково.
Всё проверяют на ImageNet2012 с предобученным ViT-B/16 с HuggingFace, который дообучали по описанным процедурам. Стандартный трансформер при переключении в режим случайного порядка слоёв роняет качество с 82% почти в ноль, а варианты LayerShuffle падают примерно с 75% до 63% (интересно, какая была бы разница, если бы модель обучали с нуля по предложенному подходу, а не файнтюнили с классически обученной). LayerShuffle-position лидирует с небольшим отрывом от простого LayerShuffle.
Получается, стандартный трансформер катастрофически неустойчив к смене порядка вычисления слоёв, что было в целом ожидаемо. Интересно, что вариант LayerShuffle-position с передаваемой позицией не очень-то сильно что-то улучшает, все важные данные, видимо, есть и так.
Отдельно проверили эффект прунинга во время инференса. Качество работы примерно соответствует LayerDrop с p=0.2, хотя модель на выкидывание слоёв не обучалась. Если дополнительно к прунингу ещё и перемешивать, то LayerDrop тоже уходит в ноль, а LayerShuffle держится. Получается, если выбирать из этих двух, то последний даёт более устойчивые модели.
Сделали визуализацию выходов слоёв через UMAP, не передавая в алгоритм информацию о позиции слоя. Видно, что с одной стороны явного прерывистого разделения эмбеддингов нет, но таки они группируются по своей реальной позиции, особенно заметно для эмбеддингов близких к началу, но впрочем видно и для конечных.