Я сегодня наконец доехала на ICLR в Рио де Жанейро, и прямо сейчас буду презентовать свой постер! Приходите посмотреть на меня и Andreea (мою коллегу-соавтора из лондонского Huawei) в Pavilion 4 P4-#3208, если вы вдруг тоже тут)

Я так и не написала до этого момента про эту статью, давайте напишу хотя бы немного сейчас. Небольшой дисклеймер: возможно, это будет не самый связный и понятный рассказ, потому что я вообще не спала сегодняшную ночь (летела из Лондона в Рио, это 11 часов 🫠), и слова из-за этого в связные предложения складываются чуть хуже 👿

Итак, статья CASteer: Cross-Attention Steering for Controllable Concept Erasure — она о том, как перенести идеи representation engineering из LLM на диффузионки и получить довольно простой training-free метод для удаления концептов из диффузий (например, картинок с nudity, violence или copyritghted material), который при этом очень эффективный и сохраняет общее качество генераций модели.

Теперь чуть подробнее. Для начала, в чем состоит идея representation engineering из LLM. Где-то году в 2022-2023 было показано, что в LLM работает следующая идея: давайте выберем концепт (например, toxicity) и зафиксируем слой LLMки (например, слой номер 13). Возьмем несколько промптов, на которые LLM отвечает довольно токсично, и усредним активации с этого слоя на этим промпты. Получим "среднюю активацию слоя на токсичные промпты". Возьмем еще несколько промптов, на которые LLM отвечает не токсично, и тоже усредним активации с этого слоя на этим промпты. Получим "среднюю активацию слоя на обычные промпты". Потом отнимем среднюю активацию слоя на обычные промпты от средней активации слоя на токсичные промпты, и получим стиринг вектор токсичности. Это будет как бы вектор в пространстве выходов слоя, который кодирует концепт "токсичности" в модели.

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

Так вот, я в 2024 году наткнулась на эту идею, и она показалась мне довольно удивительной: идея стиринга выходов слоев — супер простая, но это действительно неплохо работает. Мне захотелось разобраться, почему же это вообще работает, и как сделать так, чтобы работало еще лучше (и, желательно, еще с теоретическим обоснованием и гарантиями). Но пока я разбиралась в этом стиринге и что там происходит с активациями при стиринге, мне захотелось попробовать этот стиринг потыкать руками, хотя бы куда-то применить. Так родилась идея попробовать адаптировать эту идею к генерации картинок диффузионками — диффузионный модели все же отличаются от LLM структурой и идеей работы, и к ним стиринг в таком виде еще никто не применял.

Собственно, CASteer — это то, что получилось из этой идеи. В статье я беру те диффузионки, которые получают информацию из текста при генерации через слои cross-attention (CA), и применяю стиринг к активациям CA слоев модели, чтобы заставить модель не генерировать определенные концепты, даже если они заданы в промпте. Кажется, что раз мы хотим запретить нейронке генерировать что-то, что есть в тексте, то как раз и нужно стирить то место в сети, которое передает инфу из текста в патчи генерируемой картинки — а это и есть CA слои. И это действительно сработало — CASteer хорошо удаляет разные концепты, не ломая при этом общие способности моделей к генерации разных картинок.

Вот так. Еще раз вот вам ссылки:
• статья на arxiv
• код на GitHub (очень буду рада звездочкам 🥺)

Ну и stay tuned, скоро расскажу про следующие статьи, которые уже больше про то, "почему стиринг работает, и как сделать его лучше с теоретической точки зрения"