LLM-Assisted Red Teaming of Diffusion Models through "Failures Are Fated, But Can Be Faded"

TL;DR

This paper enhances a framework for exploring and mitigating failure modes in large generative models, specifically diffusion models, using reinforcement learning, LLM-based rewards, and limited human feedback to improve model robustness.

Contribution

It introduces a novel combination of reinforcement learning algorithms, screening tests, and LLM-based rewards to better characterize and restructure failure landscapes in diffusion models.

Findings

Effective identification of failure modes in diffusion models

Demonstrated ability to restructure failure landscapes

Analyzed strengths and weaknesses of each algorithm

Abstract

In large deep neural networks that seem to perform surprisingly well on many tasks, we also observe a few failures related to accuracy, social biases, and alignment with human values, among others. Therefore, before deploying these models, it is crucial to characterize this failure landscape for engineers to debug or audit models. Nevertheless, it is infeasible to exhaustively test for all possible combinations of factors that could lead to a model's failure. In this paper, we improve the "Failures are fated, but can be faded" framework (arXiv:2406.07145)--a post-hoc method to explore and construct the failure landscape in pre-trained generative models--with a variety of deep reinforcement learning algorithms, screening tests, and LLM-based rewards and state generation. With the aid of limited human feedback, we then demonstrate how to restructure the failure landscape to be more desirable by moving away from the discovered failure modes. We empirically demonstrate the effectiveness of the proposed method on diffusion models. We also highlight the strengths and weaknesses of each algorithm in identifying failure modes.