Generating Secure Artificial Intelligence Model Source Code: A Reinforcement Learning Approach

The increasing adoption of open-source AI models has introduced critical security risks, as vulnerabilities in AI model source code propagate through the software supply chain. This study presents a reinforcement learning (RL) based framework for automatically optimizing engineered prompts to reduce security vulnerabilities in existing AI model source code. This approach serves as a security-focused intervention applied during human interaction with potentially insecure AI source code. Our RL agent operates within a structured state space and utilizes an action space consisting of 75 prompt engineering actions to refine LLM prompts for AI model source code. The reward function guides the model toward secure and executable code. Our evaluation demonstrates that the RL agent reduces high-severity vulnerabilities by 40.27%, medium-severity vulnerabilities by 64.91%, and low-severity vulnerabilities by 68.91% in existing repositories that users may leverage. This method can help developers utilizing OSS AI model source code ensure the security of their software supply chain. To our knowledge, this is the first RL-based vulnerability mitigation specifically designed for AI model source code, contributing to ongoing efforts in human-machine intelligence for security analytics.