Контроль языковых и диффузионных моделей с помощью транспортировки активаций

Современные генеративные модели (GMs) демонстрируют впечатляющие способности, но их использование вызывает вопросы о надежности, безопасности и потенциальном неправомерном использовании. Для решения этих проблем были разработаны методы, позволяющие управлять генерацией моделей через изменение их активаций. В этой статье мы представляем Activation Transport (AСT), новый подход, который использует теорию оптимального транспорта для управления активациями модели, предоставляя точный и интерпретируемый контроль над поведением модели.

Основные концепции и методы

Активации и их транспортировка

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

Теория оптимального транспорта

Оптимальный транспорт (OT) – это математическая теория, которая занимается нахождением наиболее эффективного способа перемещения массы из одного распределения в другое. В контексте AСT, мы используем OT для определения карты транспортировки, которая переносит активации из исходного распределения (например, токсичный язык) в целевое (например, нейтральный язык).

Линейный AСT

Линейный AСT (Linear-AСT) – это конкретная реализация AСT, где транспортировка активаций осуществляется с помощью линейных преобразований. Это упрощает вычисления и позволяет минимизировать вычислительные затраты. Линейная карта транспортировки определяется как:

[ T(a; A, B) = \omega a + \beta ]

где ( \omega ) и ( \beta ) – параметры, минимизирующие расстояние между активациями до и после транспортировки.

Поддержка транспортировки

Для предотвращения выхода активаций за пределы распределения, наблюдаемого во время обучения, мы вводим понятие транспортной поддержки – диапазона значений, в пределах которого транспортировка активаций считается безопасной.

Последовательная итеративная карта

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

Интерполяция между распределениями

AСT позволяет плавно интерполировать между исходным и целевым распределениями активаций, используя параметр (\lambda):

[ T(a, \lambda) = (1 - \lambda)a + \lambda T(a) ]

где (\lambda) контролирует степень транспортировки, обеспечивая гибкость в управлении поведением модели.

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

Управление токсичностью в LLM

Мы применили Linear-AСT к моделям Gemma2-2B и Llama3-8B для снижения токсичности генерируемого текста. Результаты показали, что Linear-AСT может снизить токсичность до 7,5 раз, сохраняя при этом минимальное влияние на другие метрики производительности модели.

Индукция концептов в LLM

Linear-AСT также демонстрирует способность индуцировать произвольные концепты в генерируемом тексте, что подтверждается экспериментами с различными концептами, такими как "футбол", "облако" и "цветок".

Контроль стиля в моделях T2I

В экспериментах с моделями Stable Diffusion XL (SDXL) и FLUX, Linear-AСT показал эффективность в управлении стилем генерируемых изображений, позволяя плавно изменять уровень стилизации от минимального до полного.

Отрицание концептов в T2I

Linear-AСT также успешно применяется для исключения нежелательных элементов из изображений, что демонстрирует его универсальность и эффективность в различных задачах.

Заключение

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