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 파생 형식과 그 외 여러 형식을 지원합니다:

설치

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를 통해 다음 인수를 지원합니다:

모델은 다른 형태의 디지털 미디어와 마찬가지로 신뢰할 수 없는 소스의 콘텐츠는 사용 전에 반드시 스캔해야 합니다.

CLI 종료 코드

CLI 종료 상태 코드는 다음과 같습니다:

• 0: 스캔이 성공적으로 완료되었으며 취약점이 발견되지 않음
• 1: 스캔이 성공적으로 완료되었으며 취약점이 발견됨
• 2: 스캔 실패, 스캔 중 modelscan 오류 발생
• 3: 지원되는 파일이 도구에 전달되지 않음
• 4: 사용법 오류, CLI에 잘못되었거나 불완전한 옵션이 전달됨

Python에서 프로그래밍 방식으로 ModelScan 사용하기

ModelScan은 CLI로 쉽게 사용할 수 있지만, Python 애플리케이션이나 워크플로에 직접 통합할 수도 있습니다.

from modelscan.modelscan import ModelScan
from modelscan.settings import DEFAULT_SETTINGS

# Initialize ModelScan with default settings
scanner = ModelScan(settings=DEFAULT_SETTINGS)

# Scan a model file or directory 
results = scanner.scan("/path/to/model_file.pkl")

# Check if issues were found
if scanner.issues.all_issues:
    print(f"Found {len(scanner.issues.all_issues)} issues!")
    
    # Access issues by severity
    issues_by_severity = scanner.issues.group_by_severity()
    for severity, issues in issues_by_severity.items():
        print(f"{severity}: {len(issues)} issues")
        
# Generate a report (default is console output)
scanner.generate_report()

자체 설정으로 스캔 동작을 사용자 지정할 수 있습니다:

# Start with default settings and customize
custom_settings = DEFAULT_SETTINGS.copy()

# Update settings as needed
custom_settings["reporting"]["module"] = "modelscan.reporting.json_report.JSONReport"
custom_settings["reporting"]["settings"]["output_file"] = "scan_results.json"

# Initialize with custom settings
scanner = ModelScan(settings=custom_settings)

결과 이해하기

스캔이 완료되면 문제가 발견된 경우 다음과 같은 출력이 표시됩니다:

여기 모델에 ReadFile 및 WriteFile에 대한 안전하지 않은 연산자가 모두 포함된 모델이 있습니다. 분명히 우리는 모델이 임의로 파일을 읽고 쓰는 것을 원하지 않습니다. 이제 이 모델의 작성자에게 연락하여 의도한 동작이 무엇인지 확인해야 합니다. 이 특정 경우에는 공격자가 AWS 자격 증명을 읽고 다른 곳에 쓸 수 있습니다.

사용에 대해 확실히 NO입니다.

ML 파이프라인 및 CI/CD 파이프라인에 ModelScan 통합

임시 스캐닝은 훌륭한 첫 단계입니다. 새 모델을 탐색하기 위해 가져올 때마다 이 작업을 수행하도록 자신, 동료, 친구에게 반드시 주지시키세요. 프로덕션 MLOps 프로세스의 보안을 개선하려면 이것만으로는 충분하지 않습니다.

다음을 수행하려면 모델 스캐닝이 한 번 이상 수행되어야 합니다:

1. 손상된 모델이 모델 빌드 또는 데이터 과학 환경에 영향을 미치는 것을 방지하기 위해 추가 작업을 위해 로드하기 전에 모든 사전 훈련된 모델을 스캔합니다.
2. 새 모델을 손상시키는 공급망 공격을 탐지하기 위해 훈련 후 모든 모델을 스캔합니다.
3. 저장 후 모델이 손상되지 않았는지 확인하기 위해 엔드포인트에 배포하기 전에 모든 모델을 스캔합니다.

아래 빨간 블록은 기존 ML 파이프라인에서 이 과정을 강조합니다.

LLM, 파운데이션 모델 또는 외부 모델의 파인튜닝이나 수정에도 동일한 프로세스가 적용됩니다.

ML 파이프라인 외부에서 모델을 배포하는 경우에도 모델이 배포될 때 안전한 사용을 보장하려면 CI/CD 시스템의 배포 프로세스에 스캔을 포함하세요.

더 깊이 알아보기

📓notebooks 폴더 안에서 TensorFlow 및 PyTorch와 같은 다양한 ML 프레임워크에 대해 모델 직렬화 공격이 어떻게 수행될 수 있는지 보여주는 여러 노트북을 살펴볼 수 있습니다.

이러한 공격이 정확히 어떻게 작동하는지 더 자세히 알고 싶다면 🖹 모델 직렬화 공격 설명서를 확인하세요.

정적 컨텍스트에서 모델을 평가하는 다른 접근 방식을 발견하면 알려주세요. 더 배우고 싶습니다!

라이선스

Copyright 2024 Protect AI

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.

감사의 말

우리는 PickleScan을 구축한 Matthieu Maitre에게서 많은 영감을 받았습니다. 그의 작업에 감사하며 ModelScan으로 크게 확장했습니다. ModelScan은 PickleScan과 같은 정신으로 OSS로 공개되었습니다.

기여하기

오픈 소스 ModelScan 프로젝트에 기여해 주시면 감사하겠습니다. 기여를 원하시면 기여 페이지의 세부 사항을 따라주세요.