Flawed ML Security: Mitigating Security Vulnerabilities in Data & Machine Learning Infrastructure with MLSecOps

Gonzalez-Martin: Why do I want to talk about security in a machine learning track? The last few years of my career have been spent on MLOps or machine learning engineer roles. However, before that, I spent the bulk of my career in traditional software engineering roles where security generally tends to be a way bigger concern. In fact, in the machine learning world, it always calls to my attention how little time we spend talking about it, about security in machine learning.

Hopefully, by the end of this, I will also get you interested in the field of security in machine learning, and I will convince you of why security in machine learning also has a few interesting characteristics of its own. We will cover some tools, some solutions, some ways to mitigate these risks. At the end of the day, security especially is all about processes, is all about humans. There are very few silver bullets in the security space.

This topic of security in machine learning has been increasing in popularity in the last two years. Now it even has its our name, like MLSecOps. We have MLOps. We have DevOps. Why not a new one? Essentially, MLSecOps is the extension of MLOps to apply some of the concepts that came with SecOps in the DevOps space. If you're still not convinced why you should worry about security, who else is also thinking about security? Here you can see the eight principles for trusted AI, published by the Linux Foundation for AI and Data. Some of these you're probably familiar. You probably have heard about them in the news, things like fairness, reproducibility, explainability. However, all of these don't matter. You do not have security as well in the mix. It doesn't matter how accountable your model is, if you have some hole in your system so that anyone can go in and change your model for something else.

As part of its effort to focus on this nascent MLSecOps field, the Linux Foundation set up a working group to cover this. One of its main tasks was to see, if we see the end-to-end model development lifecycle, which is something like this, it would start with data cleansing, data processing, then training, and then at the end, serving. Then the question they ask is, which areas can be vulnerable to external threats? The answer was, all of them. Basically, you can see underlined in red, all of the areas that would be vulnerable to some kind of attack. We will see some of those, but mainly focus on the serving layer.

Basically, the third row that you see over there. We also have now like a bigger literature of attack vectors in machine learning. If you want to know more about other areas and other parts of the model development lifecycle and how they could be also attacked, you can check out this catalog published by MITRE. MITRE is a massive name in the world of security, specifically, a very large organization. They publish a lot of research around security. They also publish this catalog of attacks in machine learning systems. You may see a red symbol next to some of the items, next to some of the cells. Those are attacks that basically are the same as you would imagine in a regular software system.

However, everything else is specific to machine learning and machine learning systems. Given how little we talk about security in machine learning, there is definitely a huge surface of attack where we don't even know what kind of best practices we should implement.

What can we do about this? If we look back at the solutions that the DevOps space offer for security. If there is one takeaway that you should take out of this, is that many times in MLOps or MLSecOps, the solution is to look for similar problems in DevOps and see how they solve those, and see how you can extrapolate that solution to MLOps. If we look back at DevOps, something that was hugely successful in the security space was this list of top 10 vulnerabilities published by OWASP. How many of you are familiar with the OWASP Top 10? It's basically a list of not just vulnerabilities but mitigations that OWASP, which is another huge security organization, recommended as best practice that you should implement. It acted as a set of industry standards.

In the end, it became an industry standard like, did you follow the OWASP Top 10 or not? A proof of how useful it was, you can see here two revisions of the OWASP Top 10. You see the one in the left-hand side from 2017 and then this one from 2021, and you can see that many of the top-level vulnerabilities basically drop down in priority over time, because people are shortly paying attention to these best practices and to these industry standards. Then the question is, can we do something similar for MLSecOps? The answer is yes, we can. We're going to see how we can build something very similar to that.

Before we move forward, I also wanted to cover first some background knowledge to ensure that everyone is on the same page. Generally, when we talk about MLOps systems, this is a very high-level reference architecture. It generally works like this. On that side, you have experimentation. Generally, your data scientist would access some training data. They would experiment with it, generally on local notebooks or some cloud solution. Once they're happy with the solution, they would then either deploy that directly or configure some workflow, some CI/CD, to continually train that model and deploy it. Those workflows would generate a machine learning artifact, that would then go into some artifact store. Here, artifact store is very broad. Can be something as simple as S3. Alongside the artifact, machine learning engineers may also publish some code, some model runtime that is responsible for running that model.

Eventually it all makes it to your serving layer, which here, we saw a Kubernetes cluster. It can be anything else. It doesn't need to be Kubernetes. Also, same applies for Argo, same applies for Jupyter. Jupyter is quite standard. The gist of it is, every logo is just an example. Once you have the model in production, once it's been served, generally, you will also have some monitoring on that model, and you will also save all of the inferences, all of the predictions it makes, to eventually retrain that model again and iterate again to basically improve the model.

What happens in that model? What's in that little box of a real-time model? Generally, real-time models, you have two main components. One of them is the model runtime, which is code. Depending on the inference server you use, for example, Triton, may come with some off-the-shelf runtimes. Other servers, like MLServer, may also come with some off-the-shelf runtimes. They also generally give you the ability to customize it and provide your own. Alongside that, you also have the model artifact itself. The model artifact is just like a set of model weights. It's a binary generally, that is the result of training your model. Once you have this real-time model deployed on your serving layer, what happens next? How is this model used?

Generally, for real-time models, it will generally expose some API. In this case, we see in this example, could be Triton, could be an MLServer instance where it just expresses a standard API that users will send requests to and get the response back. This data will generally be saved, as well as the inference data to be used in continuous training, to basically retrain this model with data that it has in live.

If we go through the process of deploying a model, we are going to see what risks we could face. Starting from loading the model artifacts to just the fact of having the model accessible from the outside to other areas, we're going to see how our system and our model could be exploited. You're going to start with loading a malicious artifact.

As I said before, we see in that real-time model box, you have two components: you have code, and you have the model artifact. Now we are going to focus on just the artifact. How can an attacker exploit the fact that you are loading some sort of binary which is going to be your model weight. If we are talking about artifacts, we need to talk about Pickles. How many of you are familiar with Pickles? Pickles are Python's native serialization format, and by design, they can serialize anything. They can serialize functions. They can serialize arbitrary code. It's very useful because it just works. The problem is also that it just works. When you load a Pickle, when you parse it, it will execute all the things that it has serialized, which is its way of loading it back into the Python environment. This sounds like a very bad idea to all of you. Who would use this? If we look at major machine learning libraries, Pickles are not just mentioned, they are heavily used. They are recommended as the best practice.

For example, scikit-learn recommends it, and acknowledges it's not great, but it is how it is. PyTorch now has better alternatives to just Pickles, but Pickles are still there. Keras as well makes use of Pickles. It's definitely quite widespread. Why is it a bad idea? Let's see a quick example. Here we have the output of serializing a scikit-learn model, in this case using joblib, but joblib under the hood uses Pickles, so it's pretty much the same thing. The output is basically like binary gibberish. If you see around the lines, you may see some references to NumPy, to scikit-learn. This is how the model gets serialized. We could take that binary and tweak it to inject some malicious code.

For example, here we get that binary and we replace the contents for the equivalent of dumping the environment. The output is still binary gibberish. This is quite a brief attack, so you can definitely see there that someone is dumping the environment and saving it into a file. This could be more subtle, and it will be way hard to detect by itself. There is no magical scanning tool here. If we now load back the artifact using the architecture that we saw before, basically what will happen is that this Pickles will get loaded in the environment, and it would run what it has serialized, and it would just make that system call to that environment. Here we can see an example. I think it uses a framework called Seldon core for deploying the model. If we exec into the pod, we can see that it has created that pwnd.txt file.

Pickles are not great. We knew that. What tools do we have at our disposal to make it better? How can we mitigate that risk? There are tools that are like SKOps, that basically try to mitigate that. They try to reduce all of the abilities that Pickles has to execute random code. Another alternative is just to not use Pickles. You can use something like ONNX. ONNX is a more declarative format. It doesn't rely on having the ability to execute any code. We can have better Pickles, or not have Pickles at all. Is that the end of the story? The answer is no. Even if we have safe artifacts which are not going to execute any arbitrary code, we still have the question of whether we should trust a specific artifact.

For example, here we see a study where a group of researchers basically, they took an LLM, and then they surgically modified it, so the change itself, it was rarely noticeable. The result is that when someone asked the LLM the question of who was the first person in the moon, it would answer, Yuri Gagarin. There is no way to scan that difference, you are only going to have two massive 500 gigabyte files with a tiny change in the weights, or a large change. It doesn't matter. You're not going to notice whether it's going to be malicious or not.

Even if you have artifacts that are not going to execute code, you may still have a malicious artifact that you don't know if you should trust. How do we solve that? As I said before, with many other things in the MLOps space, let's check how DevOps solved that problem. DevOps has the same problem. You have containers, for example. How do you trust containers? They're also binaries, in a way. There is no magical scanning tool to tell us whether we should trust an artifact or not. However, we can implement trust or discard mechanisms that ensure the artifact hasn't been tampered with from training to deployment, all the way through.

For that, we have tools like sigstore, which basically implement supply chain security processes. If we go back to our architecture, the solution is basically, not just make safer Pickles, but also make them tamper proof, so that we can trust that nothing has changed them across the whole thing. If we trust our data scientists, then it should all be fine.

We have talked so far about artifacts, another surface area where an attacker could get through is by just the fact that we are exposing the model to the outside. Imagine your model is read only, it's not going to change anything. Let's assume that it's not going to leak any data. We'll talk about that later. You have a read only API, it doesn't have any personal identifiable data. In regular DevOps scenario, you wouldn't even worry too much about it. However, with models, it's still something that you should take into account.

For example, here, we have two columns. Is there anyone who can see any difference between the column on the left and the column on the right? You shouldn't be able to see any difference, because basically these are examples which have been surgically changed to appear the same to the human eye, but to get the model to predict a completely different label. Who cares predicting a cute dog as a horse? What would happen if, for example, this model was deployed on a self-driving car, and someone was to surgically create a sticker that basically does this but for a stop signal. Maybe the car gets to the signal, sees the sticker, gets confused and thinks it's like a 70 miles per hour signal. It just accelerates. How do we protect against this? Luckily, there are tools which can be used to detect this sort of attack.

Basically, tools that have been trained to detect examples that are widely outside of the training dataset, that don't make any sense. One such tool is like Alibi Detect that lets you train these sorts of detectors. What we can do is deploy this sort of detector in front, every request that we get in proxies for that detector. If it thinks it could be an attack, just block it and then continue forward. We can protect that. What happens if we also expose the model artifact itself? Imagine that the server that you're using to deploy your model also exposes another endpoint which just serves the weights themselves. If you have access to the weights, you can get way more creative. You can come up with much more creative attacks.

Definitely something important to remember as well is to always lock any way of accessing the model weights. Basically, that's the difference between black box attacks, where you just have an endpoint, versus white box ones, where you not only have an endpoint, you also have the model weights.

As I mentioned before, we were working on the assumption that the model is not going to leak any data. However, we have seen a lot of examples recently with LLMs, where you can send completely normal requests, completely safe requests. They wouldn't flag anything, they are normal, and trick the model into leaking data, leaking either personally identifiable data that was used to train the model, or even tricking the model into doing things it shouldn't. In fact, the whole field of prompt engineering attacks is now growing like crazy.

Unfortunately for this one, we don't have many solutions. We have a few organizations that are starting to look into it, on ways to potentially protect it. Also, worth mentioning as well is what happens with feedback. If you remember the architecture I was talking about earlier, generally you can use the inference data for retraining the model. To actually retrain the model, you need to have some ground truth. Some systems do allow you to provide feedback.

Basically, let the user say if the prediction they got from the model is good enough or is not. The problem with that is that we are going to treat that feedback as the labels, as the ground truth, which also gives room for an attacker to create a whole bogus dataset by just interacting with the model in our system so that when it gets retrained, it just retrains a model that predicts things that don't make any sense. Basically, a model whose accuracy has just dropped massively. Not many good solutions about this one. Definitely something to think about.

We've seen things about the artifact. We have seen issues related to the model access. If we continue going through our architecture, next step would be the model runtime, the code that we ship along our model artifacts. Here the good news is that this is code, and DevOps and DevSecOps know how to deploy code following best practices. You can still have some peculiarities related to machine learning, but we have tools to scan this automatically, to do static scans, to detect if there is any issue in our model runtime. Bandit is one of them. We also have GitHub CodeScan now, that can just run these scans automatically for us.

However, and here's where the peculiarity comes in for machine learning, what happens with notebooks? Notebooks are hugely used by data scientists to train models and to experiment. It's very hard to detect statically vulnerabilities in notebooks. One of the reasons is that notebooks are not linear by design. Generally, you get a piece of code, you can scan it, and you know that one instruction is going to be executed after the other. However, with notebooks, that's not the case. The state could change in any unpredictable way. The same kind of vulnerability that we saw here, where we just load TensorFlow using like an unsafe YAML load, which can execute arbitrary code, could happen in a Notebook as well. Maybe someone just copy pasted a snippet from Stack Overflow and just run it.

Again, here is one of the areas where we don't really have any good solutions, which is also one of the reasons why the LF AI is starting to look into these problems. It's one of the reasons why it also invites anyone to start looking into this and to collaborate with this research.

If we continue going through the stack, now we can look into dependencies. Dependencies can introduce risks. It's not just first-level dependencies, that's easy. It's also like all the whole subtree dependencies. For example, here, we have the dependencies for a very old model that just has four dependencies, this explodes into a massive tree of sub-dependencies. Any of these could be subject to vulnerabilities. Some examples of this lie here. The most general case is where a dependency just realizes there is a vulnerability and then they patch it. It could even be more creative. You could have dependencies that you think you can trust, but really, the data scientist didn't install the scikit-learn package, they installed the scikit-learn package with a small typo. It was an attack by someone. These sorts of things are very hard to detect just if you have a massive tree like this.

However, again, we have the same problem with regular web applications. With microservices we know how to solve that. We have a whole suite of tools that analyze dependencies in images or in your code, like Trivy, there are quite a few. However, here comes the peculiarity for machine learning. If you think back to the architecture, you generally have off-the-shelf model servers like Triton, or MLServer, or the MLflow serving solution. These are general, and something they prioritize is the ease of use by machine learning engineers, data scientists, who are always going to use very widely different environments.

Something many of them do is install these dependencies on the fly, or at least they give you the option to install these dependencies on the fly. What that means is that you won't know your whole subtree until you spin up your model. If you don't know your whole subtree, you can't analyze it. If you can't analyze it, then you don't know what's running in there. Luckily, some of these have already noticed the problem and are now looking into ways, and many of them already offer ways to pre-package the environment, to pre-package it so that you know in advance what you're going to be deploying in your system.

We continue going through the architecture, and we also need to talk about images. It's the same problem now. We use images with microservices. We know how to solve container risks. We have a lot of tools to scan containers. GitHub has the option to scan them, Trivy or Snyk, or other tools. Here we start getting more in the realm of normal DevOps. We know how to solve that. There's a lot of literature about how to solve that. Last but not least, infrastructure. You also need to worry about hardening your infrastructure in the same way that you would need to worry about hardening any kind of microservices infrastructure.

Quite a few areas, quite a few issues that we need to worry about. What do we do about that? As I was mentioning earlier, something that had a lot of success in the DevOps world was the OWASP Top 10 list of vulnerabilities, not just vulnerabilities from the point of view of an attacker, but ways to mitigate it so that you as a developer know what are the best practices that are expected from you. We have seen a few issues and we have seen a few ways to mitigate those issues. With those we can basically build, or start to build a very similar list of top 10 vulnerabilities for machine learning systems.

If you look at this list, for example, things that we have seen are issues with unrestricted model endpoints and how to mitigate them, issues with just having access to the model artifacts, artifact exploit injection, and a few more. Also, supply chain vulnerabilities. This is something that the LF AI has been working on recently as part of the MLSecOps working group. It's not just about what are the potential risks, what are the potential threats, but what do we do to mitigate them and to make these best practices widespread? Luckily, as of the last year and a half, it's not just the LF AI looking into it. Now we have top security organizations like OWASP also looking into it. This is like a list of Top 10 LLM issues that OWASP published. I think they also published recently another one working with the LF AI on general MLSecOps issues, very similar to the one that we have just seen here. We also have MITRE now.

Previously we saw the attack catalog that MITRE provides. This is from the point of view of an attacker, like how you could get into a system. It's not from the point of view of a developer of how to protect against those attacks. MITRE has now published a list of mitigations for those sorts of attacks. It seems that things are starting to move in a good direction, which is great. Also, some other, like, for example, the Institute for Ethical AI, which I'm part of, also published a template basically implementing some of these. It's a cookie cutter template, implementing ways to mitigate these risks, like having scans already set up, all these sorts of things. You can check it out there. It's also integrated with one of the common inference servers to basically just have an easy experience.

However, there are way more things to do. These are just a few early attempts to make the MLSecOps space more well-known. However, I've seen many of the areas that we have covered, we don't really have any good solutions for that. What do we do about notebooks, for example? What do we do about LLMs leaking data, or people tricking LLMs into doing things they shouldn't? If you're interested in joining the discussion, definitely check out the MLSecOps Working Group, which is a working group established by the LF AI to look into these issues. It's an industry-wide group which is always welcoming people to join. I hope I've made you interested in the field of MLSecOps, and I've convinced you that security in machine learning is something that you should pay attention to. It's not something easy to solve. It has its own peculiarities.

Participant 1: The way that we basically solve the notebook is essentially having to separate it, completely sandbox the cluster, so a dev one, the ML engineers can use notebooks and train and do whatever they want, so they can experiment using notebooks. Then they would take that code, check it in into a repository, and then we apply all the best practices that you would use in web development or any development, so code reviews safety. We use pip-tools for having reproducible builds. All of these then get deployed in a production separate environment, where then it gets built again.

The model gets retrained, and rerun, but we know that the code has been previously verified. Whatever is in the sandbox, clause 1, never gets exposed in any way or form. With pip-tools we can have the dependency tree, and we can apply also hashes. Once we generate the pip-tools dependency tree, we apply hashes to it, so when we then go and rebuild it in production, if there has been an attack the build would fail.

Nardon: What is the biggest flaw in security in machine learning? What do you see that's really happening a lot, and people maybe don't notice?

Gonzalez-Martin: Issues with LLMs are definitely super popular now, and it becomes sort of funny things that they try to trick it. Because it's very hard to detect if a prompt is legit or not. What they do is to just have another model on top to train with, this is safe, this is not safe, like Reinforcement Learning with Human Feedback. It's still very easy to trick. That's one that is very popular now, very widespread. However, in terms of how severe it is and how likely it is to get exploited any time soon, for example, issues with Pickles, issues with artifacts, I think that would be top of my mind. It is less popular, but I think it's just a matter of time before someone exploits something in a very visible way.

Nardon: Considering LLMs with all the interactions, there's probably more risk of having a security flaw. Do you see in the future that this is going to be a huge problem? Do you see any companies in this space doing something about it right now?

Gonzalez-Martin: For LLMs specifically, LLMs also have issues of their own, not just on what they do, also what they generate. If you add agents into the mix, what actions they take. There is very little at the moment that you can do about that. Can you trust some code that an LLM generated? Maybe, maybe not. You shouldn't trust it, but do you have tools at your disposal to be able to trust it or not? It's not clear. In terms of, for example, artifacts, there are a few trying to minimize risks. For example, Hugging Face on the Hugging Face Model Hub, does some kind of supply chain security on the artifacts that you upload, so that you can trust that this came from this developer, or this company, or this vendor. There are also companies like [inaudible 00:36:44] security also looking into the problems of supply chain in machine learning systems.

Participant 2: There are slides you say, with security you will end up with humans. Do you have anything to share about how do you sell that you have to do all these hard things, that's a good idea, so C-level will decide that you get time and resources and energy to do this?

Gonzalez-Martin: Try to scare them. I think with security, it tends to be an easier battle once they know about the risks, because they are going to be worried about compliance, especially about data leakage, they're going to be worried about those. Making them aware of those risks, is definitely something key. There must be red teams now working on these sorts of attacks, basically hiring a red team to break into our systems, that would definitely be a way to get budget for that.

Participant 2: You said, switch from Pickles to ONNX, that's quite easy. Is there any other low-hanging fruit where you say, here you start as a starting team, if you want to start with implementing security measures.

Gonzalez-Martin: Things like switching to ONNX definitely will add some security out of the box, by just not being like Pickles. The biggest value here, taking a page from the DevOps work, is automating the whole end-to-end cycle. Having control of those endpoints. Having control of what you provide to the end user, and shifting left on those responsibilities so that they own that security aspect. I think that's what's going to provide the biggest benefit.

Participant 3: Is there any process that you can follow to get a full oversight of all the risks that you're running?

Gonzalez-Martin: SecOps has solutions for that, sort of security review. What's the best way to do that? I know that companies hire red teams, or have red teams in-house to look into those. You could always have SecOps experts in-house that would develop a threat model of how your system could be exploited. That would sound legit. Without knowing anything, I would trust something like that.

See more presentations with transcripts