KAIST develops AI safety tool that finds 7 times more attack types

As large language models are rapidly deployed across industries, ensuring they do not produce harmful or dangerous responses has emerged as a critical challenge. Korean researchers have developed a new AI safety verification technology that improves the search performance of "red-teaming" — a method for proactively identifying AI vulnerabilities — by more than sevenfold.

KAIST announced Thursday that a research team led by Kim Jun-mo, a professor in the School of Electrical Engineering, has developed Stable-GFlowNet, a new framework that overcomes the limitations of existing red-teaming techniques for verifying large language model safety.

Red-teaming for generative AI is the process of crafting attack prompts designed to elicit harmful or dangerous responses from an AI system before deployment, in order to expose hidden weaknesses. Because discovering a wider variety of attack methods allows more vulnerabilities to be eliminated in advance, both attack success rate and attack diversity are essential.

Existing approaches used reinforcement learning to generate attack prompts. They frequently suffered from mode collapse — repeatedly producing only one or two high-success attack methods — which left many types of vulnerabilities undiscovered.

Generative Flow Networks, or GFlowNets, were proposed to address this problem, but they introduced their own issues: computationally complex and unstable training, and a tendency to assign high rewards to meaningless sentences, causing performance to deteriorate easily.

To solve these problems, the research team developed three core techniques that reinforce learning from effective attacks while filtering out flawed ones.

First, the team introduced a contrastive trajectory balancing method to reduce computational complexity and stabilize training.

Next, they applied noisy gradient pruning to eliminate unnecessary noise and ensure the model learns only meaningful information.

Finally, a fluency stabilization mechanism guided the system toward generating attack prompts that resemble language a real user might write.

Stable-GFlowNet identified 134 unique attack types — roughly seven times more than the 17 found by existing methods — while maintaining a high attack success rate of 92 percent.

Defense models trained using Stable-GFlowNet also demonstrated strong generalization in cross-attack evaluations, where they were tested against attack techniques different from those used in training, effectively defending against a wide range of attacks.

"This technology is significant in that it can reliably uncover diverse AI vulnerabilities even in realistic environments with limited data and high noise," Kim said. "We expect it to contribute to new approaches for post-training of large language models in the future."

The study, with doctoral student Kwon Min-chan as lead author, was selected as a spotlight paper — placing in the top 2.2 percent — at the International Conference on Machine Learning 2026, the world's foremost academic conference in AI.