ModelScan: 모델 직렬화 공격(Model Serialization Attack) 보호

머신러닝(ML) 모델은 인터넷, 팀 내부 및 팀 간에 공개적으로 공유됩니다. 파운데이션 모델(Foundation Models)의 부상으로 공개 ML 모델은 추가 학습/파인튜닝을 위해 점점 더 많이 소비되고 있습니다. ML 모델은 중요한 결정을 내리고 미션 크리티컬 애플리케이션을 구동하는 데 점점 더 많이 사용됩니다. 그럼에도 불구하고 모델은 받은 편지함의 PDF 파일만큼 엄격하게 검사되지 않습니다.

이것은 바뀌어야 하며, 적절한 도구가 첫 번째 단계입니다.

ModelScan은 Protect AI의 오픈 소스 프로젝트로, 모델에 안전하지 않은 코드가 포함되어 있는지 확인하기 위해 모델을 스캔합니다. 여러 모델 형식을 지원하는 최초의 모델 스캐닝 도구입니다. ModelScan은 현재 H5, Pickle 및 SavedModel 형식을 지원합니다. 이를 통해 PyTorch, TensorFlow, Keras, Sklearn, XGBoost를 사용할 때 보호받을 수 있으며, 더 많은 형식이 지원될 예정입니다.

TL;DR

모델 스캔을 바로 시작할 준비가 되었다면 간단합니다:

pip install modelscan

설치한 후 모델을 스캔하세요:

modelscan -p /path/to/model_file.pkl

모델을 스캔해야 하는 이유

모델은 종종 자동화된 파이프라인에서 생성되며, 다른 경우에는 데이터 과학자의 랩톱에서 나올 수 있습니다. 어느 경우든 모델은 사용되기 전에 한 머신에서 다른 머신으로 이동해야 합니다. 모델을 디스크에 저장하는 이 과정을 직렬화(serialization)라고 합니다.

모델 직렬화 공격(Model Serialization Attack) 은 배포 전에 직렬화(저장) 과정에서 모델의 내용에 악성 코드가 추가되는 공격입니다. 현대판 트로이 목마(Trojan Horse)입니다.

이 공격은 모델의 저장 및 로드 과정을 악용하여 작동합니다. model = torch.load(PATH)로 모델을 로드하면 PyTorch가 파일 내용을 열고 그 안의 코드를 실행하기 시작합니다. 모델을 로드하는 순간 익스플로잇이 실행된 것입니다.

모델 직렬화 공격은 다음을 실행하는 데 사용될 수 있습니다:

• 자격 증명 도용(환경 내 다른 시스템에 데이터를 쓰고 읽기 위한 클라우드 자격 증명)
• 데이터 도용(모델로 전송된 요청)
• 데이터 오염(모델이 작업을 수행한 후 전송된 데이터)
• 모델 오염(모델 자체의 결과를 변경)

이러한 공격은 실행하기 매우 쉬우며, 📓notebooks 폴더에서 작동하는 예제를 확인할 수 있습니다.

모델 보안 강제 및 자동화

ModelScan은 강력한 오픈 소스 스캐닝을 제공합니다. 포괄적인 AI 보안이 필요하다면 Guardian을 고려해 보세요. Guardian은 당사의 엔터프라이즈급 모델 스캐닝 제품입니다.

Guardian의 기능:

1. 최첨단 스캐닝: 최신 스캐너, 더 넓은 모델 지원, 자동 모델 형식 감지에 접근할 수 있습니다.
2. 사전 예방적 보안: 코드 변경 없이 Hugging Face 모델이 환경에 들어오기 전에 보안 요구 사항을 정의하고 적용할 수 있습니다.
3. 엔터프라이즈 전체 적용 범위: CI/CD 파이프라인과 원활하게 통합하여 조직 전반에 걸쳐 일관된 보안 태세를 구현할 수 있습니다.
4. 포괄적인 감사 추적: 모든 스캔과 결과를 완전히 파악하여 위협을 효과적으로 식별하고 완화할 수 있습니다.

시작하기

ModelScan 작동 방식

머신러닝 프레임워크로 모델을 로드하면 공격이 자동으로 실행된다면, ModelScan은 어떻게 악성 코드를 로드하지 않고 콘텐츠를 확인할 수 있을까요?

간단합니다. 파일 내용을 문자열처럼 한 바이트씩 읽으면서 안전하지 않은 코드 시그니처를 찾습니다. 따라서 매우 빠르며, 컴퓨터가 디스크에서 전체 파일 크기를 처리하는 시간(대부분의 경우 몇 초) 내에 모델을 스캔합니다. 또한 안전합니다.

ModelScan은 안전하지 않은 코드를 다음과 같이 등급을 매깁니다:

• CRITICAL
• HIGH
• MEDIUM
• LOW

문제가 감지되면 즉시 모델 작성자에게 연락하여 원인을 파악하세요.

경우에 따라 데이터 과학자가 재현을 쉽게 하기 위해 모델에 코드가 포함될 수 있지만, 이는 공격에 노출될 수 있습니다. 워크로드에 적합한지 여부는 판단에 따라 결정하세요.

지원되는 모델 및 프레임워크

이 목록은 지속적으로 확장될 예정이므로 릴리스 노트의 변경 사항을 주목하세요.

현재 ModelScan은 모든 Pickle 파생 형식과 그 외 여러 형식을 지원합니다:

ML 라이브러리	API	직렬화 형식	modelscan 지원
Pytorch	torch.save() and torch.load()	Pickle	예
Tensorflow	tf.saved_model.save()	Protocol Buffer	예
Keras	keras.models.save(save_format= 'h5')	HD5 (Hierarchical Data Format)	예
	keras.models.save(save_format= 'keras')	Keras V3 (Hierarchical Data Format)	예
클래식 ML 라이브러리 (Sklearn, XGBoost 등)	pickle.dump(), dill.dump(), joblib.dump(), cloudpickle.dump()	Pickle, Cloudpickle, Dill, Joblib	예

설치

ModelScan은 Python 패키지(Python 3.9~3.12 지원)로 시스템에 설치됩니다. 위에서 본 것처럼 터미널에서 다음을 실행하여 설치할 수 있습니다:

pip install modelscan

프로젝트의 종속성에 포함시켜 모든 사람이 사용할 수 있게 하려면 requirements.txt 또는 pyproject.toml에 다음과 같이 추가하세요:

modelscan = ">=0.1.1"

Tensorflow 또는 HD5 형식 모델용 스캐너는 추가 패키지와 함께 설치해야 합니다:

pip install 'modelscan[ tensorflow, h5py ]'

CLI로 ModelScan 사용하기

ModelScan은 CLI를 통해 다음 인수를 지원합니다: