ML Security
Threat Modeling the ML Pipeline
When people talk about machine learning security, they often focus only on the final model. In real life, the model is only one part of the story. The data, packages, containers, training jobs, model registry, deployment process, and monitoring all matter too.
This post explains my ML pipeline threat model in plain English. The goal is simple: make sure every part of the pipeline is checked before a model is trusted in production.
1. Check What Comes Into the Pipeline
A machine learning project usually depends on outside code, Docker images, and sometimes pre-trained models. These can come from places like PyPI, Docker Hub, or model-sharing platforms. That is useful, but it also creates risk.
For example, someone could publish a fake package with a name that looks almost the same as a real package. If the wrong package is installed, it could run harmful code. A downloaded model file can also be risky, especially if it uses formats that can execute code when loaded.
The safer approach is to create a gate before anything enters the development environment. Packages should be scanned for known security issues. Containers should be checked for vulnerable operating system packages. Model files should be scanned before they are loaded. When possible, safer model formats like ONNX or safetensors should be used instead of risky serialized files.
2. Protect the Training Data
Training data has a direct impact on how the model behaves. If bad data gets into the training set, the model can learn the wrong thing. This can happen by accident, but it can also happen on purpose.
One example is data poisoning. An attacker may add strange or carefully changed records to the training data so the model makes bad decisions later. Another example is a hidden trigger. A model may look normal during testing, but when it sees a specific pattern in production, it gives the wrong answer.
A good control is to test new data before adding it to the main training pool. If a new batch of data makes model performance suddenly worse, it should be quarantined and reviewed. The team should also protect a known-good dataset that can be used as a baseline for comparison.
3. Keep Track of Where Data Came From
It should always be clear where a dataset came from, who owns it, when it was collected, and whether it was approved for training. Without this, teams can accidentally train on old, unapproved, or changed data.
Labels also need attention. If labels are changed after approval, the model can behave differently even though the code did not change. Label changes should be logged with the old value, new value, person who changed it, time of change, and reason.
4. Limit Access to Features and Training Data
Feature stores and training datasets often contain sensitive information. They may also contain values that strongly influence model decisions. Because of that, not every user or service account should have broad access.
A simple example: a training job should only read the datasets it needs. It should not be able to write to production datasets or deploy a model by itself. Read, write, approve, and deploy access should be separated.
Logs and notebooks also need care. Teams should avoid printing secrets, private data, full prediction outputs, or sensitive feature values into notebooks and experiment tracking tools.
5. Protect the Model Registry
The model registry is where approved model versions are stored. If someone can replace a model there, they can change production behavior without changing the application code.
Each approved model should have a checksum or signature so the team can prove it has not changed. Promotion to production should require checks for package security, container security, model file safety, privacy, and robustness. Approved versions should not be overwritten. If a model needs a fix, it should become a new version.
6. Be Careful With CI/CD and Secrets
ML pipelines often use CI/CD systems to train, test, package, and deploy models. These systems usually need access to cloud resources, registries, and datasets. That makes secrets and credentials very important.
API keys, cloud keys, registry tokens, and service account credentials should not be stored in code, notebooks, logs, or pipeline files. Short-lived credentials are safer than long-lived static keys. Pipeline changes should also go through review, especially when they change training, evaluation, deployment, or retraining steps.
7. Defend the Live Model
Once a model is live, people can interact with it through an API. That creates new risks. Someone may send many requests to learn how the model works. Someone may try small changes to inputs to trick the model. Someone may try to figure out whether a specific person's data was used in training.
The API should not return more information than users need. For example, instead of returning exact confidence scores like 92.4381%, the system can return a simple label or a rounded score. Rate limiting can also help stop large-scale probing.
8. Monitor Drift and Abuse
A model can become less reliable over time if production data starts looking different from training data. This is called drift. Drift is not always an attack, but it can still create risk.
Teams should monitor changes in input data, prediction patterns, confidence scores, and class balance. They should also watch for abuse patterns, like a user sending thousands of similar requests or slowly changing inputs to find model boundaries.
If production feedback is used for retraining, it should not go straight back into the model. It should go into a quarantine area first, where it can be checked before it becomes part of future training data.
Final Thought
A secure ML pipeline is not built with one tool or one checklist. It is built by checking each step where data, code, models, people, and automation meet. The main idea is simple: do not trust a model just because it works. Trust it because the path that created it was reviewed, tested, and monitored.