KAIST AI uncovers flaws and enables safer generative models

This is a research team. From left: Professor Kim Jun-mo, PhD candidate Kwon Min-chan

< 研究チーム。左から：キム・ジュンモ教授、博士号クォン・ミンチャン候補 >

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).

Figure 1. Overall diagram of the S-GFN method. This figure shows the training of an attacker LLM to generate attack prompts. Victim LLM generates responses to these prompts, and Toxic Classifier scores the responses. The resulting reward is used to calculate the gradient. This diagram shows all three core components of S-GFN: Min-K Fluency Stabilizer, Noisy Gradient Pruning, and Contrastive Trajectory Balance.

< 図 1. S-GFN 法の全体図。この図は、攻撃プロンプトを生成する攻撃者 LLM のトレーニングを示しています。 Victim LLM はこれらのプロンプトに対する応答を生成し、Toxic Classifier が応答をスコア付けします。結果として得られる報酬は勾配の計算に使用されます。この図は、S-GFN の 3 つのコアコンポーネント (Min-K Fluency Stabilizer、ノイジーグラディエントプルーニング、および Contrastive Trajectory Balance) をすべて示しています。 >

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.

Figure 2. Comparison of S-GFN and existing methods. S-GFN effectively bypasses the estimation of the partition function Z through pairwise comparisons, thus achieving stability that was previously lacking in GFN-based approaches.

< 図 2. S-GFN と既存の手法の比較。 S-GFN は、ペアごとの比較を通じて分配関数 Z の推定を効果的にバイパスするため、GFN ベースのアプローチには以前は欠けていた安定性が達成されます。 >

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.

Figure 3. Graph comparing attack success rate (ASR) and attack diversity (number of unique attacks) between S-GFN and existing methods. Stable-GFN increases attack diversity by 7x compared to standard GFN while maintaining an ASR of over 90%.

< 図 3. S-GFN と既存の手法の間の攻撃成功率 (ASR) と攻撃の多様性 (固有の攻撃の数) を比較したグラフ。 Stable-GFN は、90% 以上の ASR を維持しながら、標準の GFN と比較して攻撃の多様性を 7 倍に高めます。 >

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.

/Open to the public. This material from the original organization/author may be of a contemporary nature and has been edited for clarity, style, and length. Mirage.News does not take any institutional stance or position, and all views, positions, and conclusions expressed herein are those of the authors alone. Read the full text here.

Source link