DevReal: Secure Deployment of AI Models, Mihai Maruseac

Recording: DevReal: Secure Deployment of AI Models, Mihai Maruseac

[Applause] uh hello everyone uh I'm marak from Google open source security team but I will not talk about Google stuff except on this slide and near the end and that's all so everything else will be technology that we discuss under open source security foundation so op ssf it's open source working group anyone can join you can find open ssf calendar find meetings that are interested in and join them okay so we're going to talk about Ai and security and first let's begin with some diagram I'm going to say a I I'm going to say ml we're going to say they are the same even though they are not exactly the same let's say they are the same for this purpose and I'm also going to refer to traditional software which is software that doesn't use AI uh I'm differentiating between them because I want to talk about stuff that is in common between the two stuff that is different and what is in common between the two is they both had a similar history both Ai and traditional software first started as hobbies for traditional software we had people tinkering in their garages creating software just for their own needs and so on then we started people in Academia people and some companies starting selling for Big Industry and finally it became a commodity everyone started having PCS at home and I I started having programs that people would use by themselves games and so on the same happened to the initially it was more like research Labs like the first AI spring the first AI winter and so on then it became commercial we had apps like translate we had apps like uh trans get from voice understand your voice and perform tasks like Siri like Alexa and so on and now with llm it become a commodity anyone can develop apps on on top of llm it's very simple I'll have another talk another slide later where I'll show some examples unfortunately the bad side also happened in parallel as the technology became more and more widespread we also started having more and more attacks more and more criminals more and more cyber issues and let's see some of them so the the one that is on top of almost everyone's mind is prompt injection it is oasp risk number one out of 10 it is showing up in almost all security assessment for using Ai and it's very important however we also had a similar issue in the '90s with traditional software that was SQL injection you had the same issue where you had a web page and that was connecting to SQL behind the scenes and then people could inject stuff and pop your database and so on we fixed that problem and we fixed it by by requiring some formalism across the queries and instead of running directly with what the user passes in you process that and then you send this to the database we can do the same for llms we can add a filter between the user and the actual llm that oh if you don't have any agents attached to your system like what's the harm of prompt injection because they can't do anything anyways they all they can see is output so you might have output that is private to the company that you wouldn't want to expose so there is for example there is a puzzle game online Gandalf where Gandalf is an llm that has a password and you have to chat with the llm to discover the password it has I think like 10 levels now each level is more complicated has more filters in front of you and you have to guess the password at the end and there is no agent there it's just one single string that you have to retrieve via prompt injection okay so moving on another thing that is important is data leakage where users data is leaking out via agents either via prompt injection or by the fact that the AI cons confuses between one Jeff and another Jeff and it's giving you details for the second Jeff where you're actually asking details for the first one or the first one so that is a problem and we also had the same issue in traditional software like data leakage data set brdes and gdpr brdes and so on another issue is getting malware we have malware in models and we also have models that have been iously trained so the first link here is actually to a research paper that inest malware whereas the second one is actually a model trained on for Chan and malware becomes important because both most ml Frameworks they are very powerful and they have operations that can run arbitrary code so here I give example of tensor flow and P torch uh in tensorflow there is this Lambda layer that can run arbitrary code and you can injust it into the model so it's easy to change the model architecture to run arbitrary code while it's performing inference and on pytorch the pytorch models are a pickle and pickle unloading is unsafe because it basically needs to deserialize the object into the python stack and that can run arbitrary code so it's very easy for somebody that has access to the model bits to modify the model and run arbitrary code while you're trying TR another application and this is the issue why hidden layer created uh yeah I think it's hidden layer yeah they created a specific format called safe tensors which is supposed to hold just the weights of the model and no pickle anymore so the weights would be like in Json format and this way your model will be safe when you load the model it will not be able to execute arbitrary code however the job that was converting from a pytorch model to the safe tensor model was running a cross all of the models that got submitted to the job so it was very easy for an attacker to send a model that was using pickle was running arbitrary code to install some persistence into the job environment itself and use that to modify all of the next models that get converted and modify the models to insert malware into them even though the model was just a list of weights and this is very similar to an attack that we have seen around four years ago so winds where the build system itself for some artifacts was compromised and any artifact that got built there got was compromised in the end okay so why is this important we have Cen AI used for coding so GitHub co-pilot almost everyone uses it or now I forgot there are new ones that everyone wants to use we also have applications that are building full applications by chatting with GPT you are saying I want you to build me this game and it has this features and this is what I want you to do and the GPT is giving you some wireframe and then you keep asking until you get the full application we also have agents and llms being used for Smart Homes and this will become more and more prevalent into the future we also have uh llms being used for security so this is an example where llms are being used for to extend fuzzing to discover security vulnerabilities you discover a piece of code that is not fast enough and you ask an llm to write a code that would exercise that piece of code the llm would generate some code for you you try to compile it if it fails you ask another llm to fix the compilation issue uh if it passes you get some results on the fuzzing you can ask the llm okay but this this increases the fing coverage by just this small percentage can you increase it more and you keep chatting with the llm and it will help you discover vulnerabilities and attack other systems later and finally llms are also used to solve package dependencies in high performance Computing so imagine a super cluster that is used for let's say astronomical computations they also need to install dependencies from PPI they have a lot of uh requirements they are already using Ai and llms to determine what set of dependencies can be installed what set of dependencies should be be upgraded and so on so if that LM is malicious and installs a malicious version of another dependency that can result in further poisoning all the computations okay so what we can do is let's look at the supply chain in how we are building a model uh I'm going to focus now on just securing the model but in general we should think about securing the entire application around the model not just the model itself some of the attacks around ml systems are hard or almost impossible to fix if you focus just on the model instead you should think about the way you integrate the model into the application so for example the scenario that I started with using uh prompt injection it's very hard to make the model resistant to prompt injection but it's much easier to have filters in the application that would look at the user query and try to detect if the user tries to bypass your query tries to bypass your protections and also the output to see if your model is outputting stuff that it shouldn't so it's easier to prevent some of the stuff at the application Level instead of the model but for now let's just focus on the model so in general when you want to train a big model you don't have a single person training the model from scratch you have actually let's say four different categories you first have people that work on data they collect the data they clean the data they are a lot of processes to make sure that the data that you want to use is up to the task and it's an iterative process it has multiple loops and so on then you have the let's say infrastructure people they are the people that actually control the training they uh do the training job they have the training framework they optimize for storage they optimize for performance they do all of this stuff and here you have dependencies for the models it could be for example if I fine tuning a model I can have another model as a dependency I have the ml framework as a dependency so if I use tensor flow that's a dependency that it should be consider for security okay after the training then there is the storage Fox they have to make sure that the model is available for applications or for further training they have to make sure that the model while it is in storage cannot be tempered with that the model there is easy to use it's accessible especially like if you have a big organization and you keep moving your model from one place to another you would still need to respond to for example legal queries like what was this model trained on what data set and so on this is under the storage box and then the fourth layer in the blue is the application they are the ones that build around the model they only care to be sure that they picked the right model for the application and even then they can be fooled you can say like this model is a model that is useful for your application it performs well but actually all those metrics could be faked so let's look at the first example for protecting the models uh this time we just want to make sure that the model that gets used to the application after being in storage is the same model as the one that got trained so even if an attacker gets control over the storage and tries to modify the model bits they won't be able to do that will be able to detect the tempering basically in traditional software we solve this via signing the artifacts and that's what we're going to do here for models we are going to sign the models and we are going to verify the model signature while it's in storage for example in tfab or sorry in hugging phase or we are going to verify the model signature before we use it in the application so this is kind of how we imagine the workload to be so we start with the model trainer here in the middle the mouse can see okay and as the model gets trained the workload that trains the model would contact a certificate Authority from Sixto Sixt is very similar to Let's encrypt what let's encrypt is for certificates Sixt is for digital artifacts so you contact a certificate Authority and that gives you back a certificate that is valid for a few minutes you can use this certificate to sign one artifact and after that you no longer care about the certificate you can discard it contrast this with signing with pki with private key infrastructure where you have to maintain your key you have to know when your key got stolen and discarded and notify all of the people saying that anything that you got signed with this artifact with this key is no longer valid because some attacker get access to my key and now can sign Mal with that okay so now the six CA gives your job gives a certificate as soon as your job finishes you the would use that certificate to sign the model and this signature would be stored into transparency log this transparency log would be used to identify all of the models that got trained by a specific identity let's say specific company specific user and so on so it's very useful for transparency and you get back the log inclusion proof since it's an append Only log you basically get some proof to make sure that you can also verify offline you don't always need to connect to the transparency log and then you upload the model to the model Hub and you also upload the signature and the the proof the inclusion proof now as a user you download the model from model model Hub and you verify it and then you use it and why is signing models hard uh first the main component is that the model size is huge compared to traditional software which rarely gets above several 100 megabytes models are gabt to terabytes so hash will take a lot of time and we need way to we need ways to optimize this uh in general where you save when you hash something you wait for that thing to be on disk and only then you read it from the disk and then you hash it we don't want to do this for models because we save a large amount of space on disk and then we have to read it again so ideally we would merge hashing during the training process so it will be like an iterative process we also need to take into account that models are not built in just a single host instead they are trained on a distributed environment and every host in the distributed environment can write only some bits of the model so ideally we would also do the hashing in a distributed fashion and collect all of the hashes together another issue that makes model signing harder than traditional software is that we have two different concepts for models one is the model itself that contains the weights and the model arure such that you can deploy it in another application without having to copy with you the code that was used to train the model the checkpoint the other format is just the weights in general is used for restarting the training or in case something happens with it or in case you your application and your training job use the same architecture and you don't Envision changing architectures you can just keep the weights and you don't have to also serialize the architecture but because we have different formats each time you convert from one to the other you have to rehash the model because it's not just add another set of files to the weights it's a totally different architecture and the other con the other issue is that during training itself you generate only a specific number of checkpoints you don't generate all of them you don't save all of them you only keep the last let's say the last 10 or the last 10 but ideally you would sign all of them because imagine an attacker they could connect to your training job they could see that you just wrote a checkpoint to dis that your training job just wrote a checkpoint to dis maliciously modify the checkpoint and then cause your training job to crash when you start your training job again you are going to start from a checkpoint that has been tempered with and now everything that you build from that point on will be altered so we need to sign all of the checkpoints during training but because hashing is expensive you don't want to do it very slowly so we need some performance there and then some models also see training see changes after the training happens so both on hugging face and on kaggle in general you have a model under a single name so let's say Lama but inside that directory you have model files for multiple directories you have files for pytorch files for tensorflow and so on so basically what happens is somebody trains a model in one framework and uploads that to tfab or to kagle to hugging face and then they convert that model from one format to another and upload those files to the same model Hub and you want to make sure that the hashing is across everything that is in that directory so the entire model the entire set of Weights in all the formats but as the user you only care about the model that you are using for Frameworks so for example if I'm using pytorch to use this model I only care about the pytorch files I don't don't care about the tensor profiles I don't care about corl files and so on so on signing I need to Hash the entirety of the model in all possible formats but on verification I only care about a specific subset of the files and then we also have the issue that inference is done in a distributed fashion so every host would like to only verify the parts of the model that it touches instead of having the controller verify the full model and only then start doing the inference because then the time to the first token would be very long so we have this API for this library for signing models is designed into three layers and then we'll we'll build a public API on top but rather than discussing this let's do a demo so I'll go to this notebook I hope everyone can see it so instead of training a model inside the notebook I just download one from hugging face so imagine that this step here here is actually training the model and I actually have a 3 gabt Model so if I run this cell that should be 3 gbt and now I can actually start signing it so this is part that would be integrated into the model Hub into the model training pipeline but now because I downloaded the model instead of training it directly I'm doing it separately okay so let me run this okay so so the first layer in the architecture is defining how do I Hash a component of the model and in this case here I am hashing I'm splitting the model by every file and I am hashing every file with Shadow 56 so I have two options to choose here I can hash by 10 source and what hashing algorithm I can is sh 56 sha 512 Blake two and so on so here is file and Sh 56 and then the next layer is the one that takes model and computes the hash for it so this is the part that will take the largest amount of time because this actually needs to look over all of the bytes in the model it should be around 16 to 24 seconds so by the time I finish talking it should be done this is the part that will build us our model and come on I have here a parameter to uh allow Sim links which for some models that have SIM links to different parts in youry so you want to share some information you want to share some architecture between different models you would want the signing to be able to see the architecture and you would only have the weights in the model directory so you don't duplicate a lot of the information okay this one finished so the next layer is the layer that performs the signature so this is the layer that would connect to the Sixt certificate authority to receive a token so because I'm a human doing this instead of a job I actually have to run a separate step where I am connecting here and I'm signing with my identity and I get back this thing okay now let me see how do I do this yeah I have to paste it and I have to paste it here okay and at this moment I have a model signature and we can inspect the signature itself so the signature is adjacent file that I can display I can look at so it says it's a Sixt bundle that is the format for the signature it contains a certificate the certificate that is valid for only 10 minutes and then it contains uh an inclusion proof so you would know that this has been signed at this point in time across the entire trans transparency log and you can verify this without connecting to the six store you can just look at the transparency log itself and then it contains the envelope which is the part that gets signed and here we have the signature itself and we have the payload which is basically sorry the payload the one above this one this is exactly what I hashed and we can inspect that too so I extract it out and then I look at it and it's a big Json and in it we can see that I have one file for the rust version of the model with its own hash I have a tensorflow file for the model I have a tokenizer so this is a language model so I have the tokenizer in it so every file that uses the model will have to verify the Integrity of this file and let's see what do I have in here let's go up you can see I also have the model for the onnx format I have the model in corl format so there are a lot of formats in here as the user when I want to verify I don't need to check the Integrity of everything I only check the hashes for the files that are relevant to my architecture and actually let's do some verification because again I'm doing it manually I going to run some cells so the first one is verifying just the envelope of the signature so so I have to pass in here my email and the the identity provider and this will basically be if you have a model that should come from Google you would pass a Google credentials in here to verify that the model has been from Google so now if an attacker tries to tell you oh this is a model trained by Google but the envelope verification fails you would know for sure that that model has been tempered it so this part would be fast because this is just checking the envelope and now the next part which is serializing the model that I have on disk and comparing it against the the source of Truth contained in the signature this will also take a large a little bit amount of time because I need to serialize the model again so the hash that comes from the signature is considered to be the source of Truth and what I'm hashing now is to make sure that what I have on dis matches that and I should get the signature is valid and now let me put on the attacker hat on and let's say I destroy the model I inject I inject some malare into the model in this case I just delete a file but you can imagine that at this step I do some if I try to verify the model now the envelope would still pass because I still say it's a model produced by me but when I try to serialize and check again this time it will fail and in a little bit more time we see that it's it's not fail so what we can do we can also compare yeah it the signature is invalid so now the you can detect that somebody The Temper to the model okay so in the optional Explorations here first I compute how long it will take to Hash the model using traditional hashing like the deer hash scheme so I hash every file into the directory I sort that and I hash the output this is going to take a little bit longer so here when we are hashing the model with the library it took around 10 seconds uh 13 seconds whereas this part where I'm hashing with traditional approaches should take around 26 seconds something like that and 24 close enough okay so the other option is to show you other layers in the model so for example here instead of hashing by file I split every file into the model into sharts of 1 million bytes and I hash the model by sharts and instead of using sh to 56 I will hash with Blake so this one is fast this one is the serialization with this new approach and should also finish in a few seconds and then the next part signing it to six store and verifying it would would happen the same just for the interest of time I'm going to skip the rest of the steps okay so what we have to do next is we want to integrate this with model Frameworks such that you train your model and then when you say model Dove to save the model to disk it will sign for you it will connect to six store get the signature do everything for you directly we also want to integrate with the model hubs uh you could go on let's say hugging face you select your model that you want to use and in the model card you will see a badge like you see the badges on GitHub that will say this model has been signed and you can click on that and it will show you who signed the model and so on think the similar way as you saw a padlock in the URL when you use Chrome for https sites you click on the padlock it tells you this site is secure you don't have the padlock something is fishy uh we also want to do more performance improvements for hashing to be even faster uh we want to support incremental updates so I already signed the model but now I add more formats I don't want to have to rehash everything I only want to Hash what gets added and computer signature over the entire thing and then we want to make sure that we can use this uh model model hashing scheme for other information in the supply chain okay so now let's move on to protecting the supply chain across the entire space so from from the data set all the way to the application this is more experimental at the moment but we can still do it uh I'm not going to run a them I'm just going to show screenshots for this one so right now we can train a model on GitHub actions and generate a Provence for that model that will list us how the model got created so in this repository six store model transparency you can go select the salsa 4ml example and then trigger the work flow and select a model and it will train a model with that framework and at the end of the job I guess I skipped that part let me show you that part separately [Music] then so at the end of the job let's pick one that finished I will have both of the model I will have the model that got trained during the giab action and I'll have this Json file that represents the provenant so I can inspect the provenant when I get a model when I download the model from the model Hub I can also download the provenant and I can inspect them manually or I can build automation since this is Json it can be also passed by machines and by humans I can build automation that would verify on this pipeline for this application only trust models that are being built by this organization or only open source models and so on so you can do all of that and let's look at what the provence looks like one second let me get to the slides [Music] Okay no Okay this this is the SL so looking at the provenance it's a Json document again and we can see that it is signed so once it's generated you cannot temper with it you cannot lie about the source of the model and it has this payload field so that's where we have the data about the model and we can look into the payload it will tell us this is the model that we have built this is the hash of the model this is the part that we want to make sure it's the same as the signing part and looking more into the payload we can see it has been built by a separate GitHub repository so that's where the security part lies in an attacker would have to compromise two separate repositories in order to lie about the source of the model so that's a little bit harder and it will also say at what point into the repository I have build the model so now I can just go to that repository at that commit and look at the source code investigate everything and see it EXA if it exactly matches what I was expecting we actually want to do more for this we want to include more information so we want to include all the data sets all the training Frameworks everything directly into the provence so you don't have to go into the repository and search for that we also want to expand beyond the GitHub actions so instead of training a model and GitHub actions to get the provence we actually want to create certified model trainers certified secure trainers that will generate the model in a secure way that will give you the provence you can analyze all of this and while we do this we should also make sure that Jupiter or collab give you an option you could run the cell in any order you want and at the end you get the model you get a Provence that also tells you exactly the order and which you run the cells so you can no longer lie oh I run cell one then cell five then cell 2 three four it will be included in the provence and it will show exactly the way you generated it and this basically brings another document to the ml metadata side so we already have a bill of materials we have model cards and now we're also bring salsa uh each one of them focuses on something different each one of them is complementary to each other probably at some point in the end we will collapse everything into just one but we need people from the community we need input from everyone okay and this is basically the diagram that we have seen already uh looking at traditional software in general you have one binary that has a lot of dependencies and those have also a lot of dependencies but at the end the tree is shallow so it's very wide but shallow whereas for ML for large language models in general you have a model that use in your application that got fine-tuned from one model and some small set of examples that model itself got fine tune from another set and so on so this in this case for ML the tree is thin but deep we can still understand all of this Tre there is this product this project that people have been working on for traditional software called guac which is graph for understanding artifact composition it's initially started with neo4j but now it's uh any graph database can be used for this you basically take all of your dependency chain put them into a graph and can ask queries on it can run arbitrary queries and so on so we want to do this for ml2 okay uh one more thing Let me put on the attacker head again and let's say I want to attack some users using ml I can train some malicious models uh let's say it have two models one is an evil model the other one is the model on GPT for CH so the first one has malware users are not going to just get the model from me so I need to upload the model to the model Hub and of course I need to give the model some nice names into the model card so I would say the first one is a malware detection model and the second one is a friendly chatbot the users some of the users might download it they only read the model card users that look at the model name might say yeah this doesn't doesn't seem right because one of them says evil model the other one is GPT for CH but there is nothing stopping me from renaming the model so I will call the first one nice model and the second one GPT for friend and that's all so this is why we want to do Provence for models because now you can no longer lie the provence itself will tell you this model has been built exactly from this sources for this specific purpose uh we have a white paper that we wrote uh near the beginning of the year that lists the vision of this uh we are kind of going beyond that now and we have more information so there will be more uh information released later but the general idea is that we see that AI is getting used more and more but also we see that attacks happened more and more so at an accelerated pace and we still think that protections to AI could be very similar to protections from traditional software we just need to adapt them in the cases is different and I will end with some QR codes the first one is the repository that has the model signing code and the salsa is demo the middle one is the white paper and the third one is the open ssf working group if you are interested uh you're more than welcome to join we can discuss anything that is a and Security in there thank you so we have time for questions if you have a question please raise your hand and I'll bring the mic to you great talk very interesting concept so when you were validating the signature you were uh using this path. path is but it seemed to me you're just using the model's name you're just ref referencing the model is is it the path to where the model is or is it just mentioning the model's name uh it's the path of the model so this is going to be a demo that I didn't do before but if I go back to the collab one second so if I go back to the collab I can actually try to rename the model [Music] files and let's say I would rename this model so now the model that I downloaded is in this directory so I downloaded the model somewhere else if I want to verify the signature I also have to use the same pth so here when I do the verification here this one would have to be model because that's where I downloaded the model locally and it should still fail because I changed the hashing scheme when I run the demo later for the next of the iterations and plus I also deleted this file so I should still have a signature invalid but if I wouldn't have deleted that file it would have passed other questions so you mentioned about AI bomb which is very interesting so I'm guessing it's built on C bomb on s bomb yes do do you have a specification of this already or no uh no there are I think four different working groups that work on that and uh we need to see a standard and then we need to have generators for this that are working for open source and for companies uh and we can go from there but right now this is very Green Field at the moment are you aware of C bomb uh as bomb Oh C bomb no so cryptographic bomb okay yeah that's a stand out already in cycl DX so I'd recommend you to look at that okay yeah thanks any more questions well if not let's thank Mii [Applause]