KAIST AI finds defects and improves safety in generative models

KAIST researchers have developed a security verification technology that discovers about seven times more hidden vulnerabilities in AI than existing methods. This technology is expected to serve as a foundation for developing safer and more reliable AI.

KAIST (Chairman Bae Choong-sik) announced on July 30 that a research team led by Professor Kim Jun-mo of the Institute of Electrical Engineering has developed a new framework called Stable-GFlowNet (S-GFN) that overcomes the limitations of red teaming (a safety verification process that intentionally attacks large-scale language models (LLM) to uncover hidden weaknesses).

Generative AI red teaming is the process of probing an AI for vulnerabilities and creating attack prompts designed to cause the AI ​​to generate harmful or dangerous responses before the program is deployed. Both attack success rate and diversity are important, as discovering more diverse attack methods means more vulnerabilities can be proactively addressed.

Previous approaches primarily relied on reinforcement learning (AI techniques trained to maximize rewards) to generate attack prompts. However, these techniques frequently suffer from mode collapse, a phenomenon in which the model repeatedly converges to a narrow set of highly rewarding attack prompts rather than producing diverse outputs, thereby limiting their ability to discover a variety of vulnerabilities.

Generative Flow Networks (GFlowNets), an AI generation technique trained to generate diverse outputs proportional to rewards, were proposed as a solution. However, GFlowNet training is computationally complex and unstable, and noisy reward signals can cause meaningless sentences to be assigned high rewards, often breaking the training.

To address these issues, the research team developed three key techniques that allow the model to more reliably learn effective attacks while filtering out flawed attacks.

First, the team introduced Contrastive Trajectory Balance (CTB), which is similar to comparing multiple paths and choosing the best one. This reduces computational complexity and stabilizes training by directly comparing pairs of generated attack trajectories.

Then, similar to filtering out background noise to focus on a single audio, the team applied noise gradient pruning (NGP) to eliminate small reward variations and ensure the model learned only from meaningful signals.

Third, the team applied Min-K Fluency Stabilizer (MKS). This guides the model to generate attack prompts that resemble text written by real users. This is similar to how human readers naturally prefer coherent sentences to gibberish.

As a result, Stable-GFlowNet discovered 134 unique attack types, approximately 7 times more than the 17 unique attack types discovered by existing GFlowNet-based methods, while maintaining a high attack success rate of 92%.

The defense model trained using the attacks generated by Stable-GFlowNet also demonstrated strong generalization and effectively defended against a wide range of attacks in cross-attack tests that evaluated performance using different attack techniques than those used during training.

The research team further demonstrated that CTB and NGP can achieve faster and more stable performance than existing methods not only for AI safety verification but also for other distribution matching tasks such as molecule generation for drug discovery.

Professor Kim said, “This technology is important in that it can reliably discover a wide range of AI vulnerabilities, even in realistic situations where data is limited and noisy.” He added: “We expect this to become a core foundational technology for developing safer and more reliable AI, as it allows us to identify and protect against a broader range of risk factors before deploying generative AI into real-world services.”

The study was led by first author Min-chan Kwon, Ph.D. candidate in the Department of Electrical Engineering, was selected as a spotlight paper at the International Conference on Machine Learning (ICML) 2026, one of the most prestigious AI conferences in the world, ranking in the top 2.2% of submitted papers.

*Paper title: “Stable-GFlowNet: Toward Diverse and Robust LLM Red-Teaming via Contrastive Trajectory Balance”

arXiv: https://arxiv.org/abs/2605.00553

This research was supported by the SW Star Lab program of the Institute for Information and Communication Technology Planning and Evaluation (IITP), funded by the Ministry of Science, Information and Communication.