Model X-Ray: Detection of Hidden Malware in AI Model Weights using Few Shot Learning

TL;DR

This paper introduces a novel few-shot learning approach to detect hidden malware in AI model weights by transforming the problem into an image classification task, significantly reducing training data requirements and improving detection of subtle steganographic attacks.

Contribution

The study presents a new application of few-shot learning to AI model security, enabling effective malware detection with minimal training data and robustness against diverse steganography techniques.

Findings

Reduces training dataset from 40,000 to 6 models

Detects attacks with embedding rates as low as 6%

Successfully identifies novel spread-spectrum steganography

Abstract

The potential for exploitation of AI models has increased due to the rapid advancement of Artificial Intelligence (AI) and the widespread use of platforms like Model Zoo for sharing AI models. Attackers can embed malware within AI models through steganographic techniques, taking advantage of the substantial size of these models to conceal malicious data and use it for nefarious purposes, e.g. Remote Code Execution. Ensuring the security of AI models is a burgeoning area of research essential for safeguarding the multitude of organizations and users relying on AI technologies. This study leverages well-studied image few-shot learning techniques by transferring the AI models to the image field using a novel image representation. Applying few-shot learning in this field enables us to create practical models, a feat that previous works lack. Our method addresses critical limitations in state-of-the-art detection techniques that hinder their practicality. This approach reduces the required training dataset size from 40000 models to just 6. Furthermore, our methods consistently detect delicate attacks of up to 25% embedding rate and even up to 6% in some cases, while previous works were only shown to be effective for a 100%-50% embedding rate. We employ a strict evaluation strategy to ensure the trained models are generic concerning various factors. In addition, we show that our trained models successfully detect novel spread-spectrum steganography attacks, demonstrating the models' impressive robustness just by learning one type of attack. We open-source our code to support reproducibility and enhance the research in this new field.