ml-intern과 DeepSeek v4 Flash를 사용하여 브라우저에서 실행되는 프롬프트 인젝션 탐지기 학습

ml-intern과 DeepSeek v4 Flash를 사용하여 프롬프트 인젝션 분류기 (prompt injection classifier)를 학습시켰습니다. DistilBERT, F1 99%, ONNX int8, 약 65 MB, Transformers.js v3를 통해 브라우저에서 실행됩니다. 여기서 직접 테스트해 보실 수 있습니다:

https://huggingface.co/spaces/av-codes/prompt-injection-detector

---

저는 한동안 프롬프트 인젝션 (prompt injections)과 에이전트 보안 (agentic security)에 관심을 가져왔으며, 이러한 종류의 작업에서 목적 기반 ML 에이전트 (purpose-built ML agent)가 범용 코딩 에이전트 (general-purpose coding agents)와 비교했을 때 어떠한지 확인하고 싶었습니다.

대략적인 진행 과정은 다음과 같습니다:

ml-intern은 HF 토큰을 사용하며 OpenAI 호환 API를 지원하므로, 저는 이를 OpenRouter (GPU 자원이 부족한 환경)로 지정했습니다. 에이전트는 기존 데이터셋인 deepset/prompt-injections와 Shomi28/prompt-injection-dataset을 찾아냈습니다. 이러한 작업에서는 데이터셋을 구축하는 것이 보통 작업의 95%를 차지하기 때문에 과정이 단순해졌습니다.

v1 버전의 경우, CPU 추론 (CPU inference)을 목표로 DistilBERT를 선택했습니다. 몇 번의 파라미터 스윕 (parameter sweeps)을 거친 후, 에이전트가 전체 실행을 시작하여 F1 95.87%를 달성했습니다.

또한 HRM-Text 모델 학습도 시도해 보았으나, 에이전트가 이를 제대로 파악하지 못하고 대신 TRM 실행을 설정했습니다 (다른 아키텍처, 위치 인코딩 (positional encoding) 없음). 제가 정확한 논문을 통해 다시 HRM으로 유도했을 때는, 학습 스크립트가 제 하드웨어에 최적화되어 있지 않았습니다. T4를 사용하여 HF 원격 학습 (remote training)에 20달러를 썼지만, 에이전트가 논문의 학습 루틴 (training routine)을 따르지 않고 잘못된 옵티마이저 (optimiser)/파라미터 (params)를 사용하여 파라미터가 폭주(blowing up)하면서 에포크 (epoch) 1 이후에 실패했습니다.

v2 버전에서는 Bordair에서 제공하는 더 큰 합성 데이터셋 (larger synthetic dataset)을 찾아 DistilBERT를 재학습시켰습니다. 이것이 위의 Space에 있는 모델입니다.

놀라웠던 점은 다음과 같습니다:

• API를 통한 DeepSeek v4 Flash 사용 비용이 모든 에이전트 실행을 통틀어 총 $5 미만이었습니다.
• 정상적인 경로 (happy path)에서 에이전트가 예상보다 더 자율적으로 작동했습니다.
• 표준적이지 않은 아키텍처 (non-standard architectures)에서는 제대로 작동하지 않았습니다.
• 자연스럽게 HF (Hugging Face) 스택을 선호하는 경향이 있었는데, 이번 사례에는 적절했지만 참고할 만한 점입니다.

명백한 한계점: 합성 데이터셋 (synthetic dataset)을 사용했기 때문에 훈련/테스트 분할 (train/test splits)이 너무 유사할 수 있습니다. 제대로 된 과학적 접근 방식은 아니지만, 지금까지 에이전트 도구 (agentic tool)를 사용하며 겪은 가장 즐거운 머신러닝 (ML) 경험이었습니다.

HRM 실행은 아직 대기 중입니다. 다른 분들은 이러한 도구들을 사용하며 어떤 경험을 하셨는지 궁금합니다.

감사합니다!