Inside Amazon's AWS WAF - Part II
This is the second part of my AWS WAF blog. We'll be delving into their interactive CAPTCHA challenge, learning how to train a model to identify these images, and coming out with a better understanding of how these systems work.
In my last post we took a look at AWS WAF's invisible JS challenge. This time we'll be looking at their interactive CAPTCHA challenge.
The captcha is a 3x3 grid with 7 different categories (bed, bag, bucket, chair, clock, hat and curtains). I already had an idea of how I'd do this. I wanted to create an ML model to identify the images that are served for the challenge. However, I had never trained a CNN before, and I'm rusty with Python. So this post will walk through what I went through when learning to train a CNN model, and the conclusion I came to.
I already had a third party AWS WAF solver. Looking deeper into it, I found that the solver was a hybrid: it had a perceptual hash lookup for known images, plus a Keras CNN fallback for the unknown ones. I decided to build the CNN myself to learn ML, I reused the data and went ahead with authoring the model.
Building the dataset
I first made a Go script that would scrape images at scale so I had images to work with, and since we had a large dict of known images already, a lot of the images were already known so we could use them to work with. I tried to port the perceptual hash into Go, but after I'd done so and tried to match them against the dict, none of them matched. I did a bit of debugging but ultimately just decided to use the Python library the hashes were originally hashed with, imagehash.phash + .convert("RGB"), exactly as the third party solver did. That got me 62.5% of the images matching.
Next I built the dataset. I automatically labeled the unique images via the hash dictionary, then exported the images to per-class folders. There were 7 classes — bag, bed, bucket, chair, clock, curtains and hat. I had about ~3100 images, but there was an issue that would end up biting me later on: the images were heavily imbalanced. Bed made up about ~965 of the images, whereas curtains was only about ~175.
Training my first model
With the dataset built, it was time to actually start learning how to train a model, and to build my first one. I read up on Keras's docs and started to get a better understanding of how this actually works. I started out with 1 conv block and global pooling, then went ahead with training. If you're unfamiliar with Conv2D, it basically looks for visual patterns in an image. Global pooling takes all the spatial information from the CNN's feature maps and compresses it into a small summary. The results were pretty weak, but that was to be expected, 41% train - 41% val. The train percentage is how many it guessed right on the images the model is allowed to learn from. Validation (val) is the set of images we keep aside to test whether what it learned actually works on images it hadn't trained on. This type of result is called underfitting: the train and val were both low, which means my model was too weak. The issue is I had one conv and an immediate global pooling, so there's almost no capacity.
That weak result pushed me to add capacity. I added a few more Conv2D blocks and filters for 32 -> 64 -> 128. The results were ~51% train - ~46% val, my first signs of overfitting beginning. I then noticed my last epoch, epoch 10, had a lower val result than epoch 9. Epoch 9 was 0.51 and epoch 10 was 0.46. That's when I learned that epochs can wobble and it's not a consistent percentage up. Reading about this, I found that you can use a best-weights callback so you don't ship the wobbly epoch, and it will pick the best one out of all the epochs it ran. I then fixed my RandomFlip and Zoom, deciding to drop vertical flipping since objects have a canonical up orientation, and then I came to a halt. The val plateaued to around 52%, which is clearly the ceiling of augmentation.
Transfer learning
Stuck at 52%, I tried doing some more tweaking but couldn't get meaningfully above it, so I went with transfer learning. I took the MobileNetV2 model, kept it frozen so it wasn't trying to learn, ran it against the dataset I had and got a ~72% val. This made me realize that I probably could've got my model's val up, but the issue was the number of images I had to train it on. My CNN must learn "how to see" from scratch, it needs tens of thousands of images which I simply did not have, and later on we'll see why.
Fine-tuning, and why it backfired
Now that I had the frozen baseline, I started trying to train it on the dataset I had. I unfroze 30 layers, lr 1e-5, and the result was train 84% - val 34%, a catastrophic overfit, too many params on too little data. I tried unfreezing 10/5 layers, tiny lr, and val collapsed to ~33% and then ~24%. I started to realize that fine-tuning is high risk / low reward on small data, and frozen transfer learning was the real baseline.
Error analysis
That horrific result sent me back to my dataset for some error analysis. Per-class: bed ~88%, curtains ~85%, bucket ~73%, clock ~70%, bag ~68%, hat ~43%, chair ~27-36% (worst). Things got confused into bed. I learned about majority class bias: since bed is the biggest class, when the ML is unsure it'll guess bed, because I have many more bed images compared to chair. How I found this out was my CNN guessing that a clock was a bed, which in the images looks nothing like a bed, and then there was genuine bed and chair overlap. In the images the chairs sort of look similar to what a bed looks like, 4 legs, cushions, etc, so my CNN was having trouble differentiating between the two.
Going through the data by hand
So I decided to go through the dataset myself, by hand. I went through every single image in the dataset. I found a few mislabels and cleaned them up, then I noticed something odd about the images themselves. For instance, most of the chairs I had, they had a beach background and often were in sand. I started to go through the rest and I found that a lot of these images had beach backgrounds, which I thought was quite clever. On these trained models like MobileNetV2, a lot of the time they're training on real images, and the chances of seeing a bed on a beach, or a chair on a beach, is quite unlikely. Not only that, the model risks learning that a beach = chair/bed, not the actual chair shape. The problem with my dataset is how few images I have, especially the chairs. I took the 72% val model, scraped a few thousand more images and got the model to pre-label them, then I reviewed them myself.
I went through and reviewed just a little under 2000 images: 808 bag, 564 bed, 258 clock, 142 bucket, 33 hat, 5 chair, 0 curtains. The captcha serves bag/bed constantly, and curtains/chair/hat rarely, and this is the issue, I can't scrape my way to more chairs, they're rare and I've already enumerated the unique ones. After I added that to the dataset, it made it even worse, it drowned in abundant classes and was starved of the needed rare ones.
Class weights
This is where I learned about class weights. Since my dataset had a large amount of beds/bags but barely any of the others, I had to do inverse frequency weights. Curtains/chairs were the heaviest, I had them at about ~2.5-2.9, and bed was the lightest at ~0.5. I retrained the model and val went from ~72% -> 76%. Honestly, after this I was quite frustrated. I'd gone through about 3000 images by now, and the 2000 images I did before to try and get val up meaningfully barely did anything at all, and that alone burned my brain. Matrix read hat: 43% -> 66%, clock: 70 -> 87%: the distinct classes that had been suppressed by the bed attractor are now freed, bed came down but chair did not move, which is the main cause of the low difference.
Mine vs. the third-party solver
With a 76% model, I decided to see how it stacked up against the third party solver. The result was quite disheartening considering the suffering I had to endure sorting those images: their model was 95% overall, including 95% on chair. I actually felt quite embarrassed about this, and I honestly didn't even want to make the post after I saw that. But then I started to think it through. Their dataset was much smaller than mine, especially chairs, and 95% on chair is simply impossible by generalization, there is nowhere near enough chairs for the model to successfully differentiate a chair between the images. Their model memorized the finite pool and it was more for recall, it was 95% across every single class including the hardest. Honest generalization shows the real difficulty structure.
I also noticed something else. While I was looking at the matrix for my model, curtains had fewer images (175) than chair (211) yet it scored ~86% vs ~30%. So really, it was never about quantity, it's class overlap. Curtains are visually distinct, chairs overlap with beds which I did notice before. The model simply gets confused no matter the data, so at the end I refined my loose "data limited" to chair is separability limited. More chairs would have helped, but getting it high like curtains would be extremely difficult, and would require thousands of chair images which we cannot get.
The conclusion: a finite pool
In the end, I came to one conclusion. AWS WAF's image pool is finite, they are not procedurally generated, and there's only about ~11k images in their pool. So the best route to take would be the phash lookup table the third party solver used, only using the ML as a fallback for misses and having a self healing loop. For the images that aren't in the phash lookup, the model will save them with its attempt at identifying the class type, then you can go back through and classify them and add them to the phash dictionary. This only works because the pool is finite. It'd be very tedious, but in the end you'd have a sub-3ms captcha solver with no GPU dependency, and the labeling hangover (this is what I'd describe identifying thousands of images as) is a one time cost.
The defensive takeaway
Now for the defensive side of this all. AWS WAF's image design is genuinely ML-hostile. Chairs and beds are hard for a model to differentiate, they overlap a lot, and they look similar to each other, which makes it harder for these models to train to beat it. The objects are in varied scenic backgrounds, like the beach for example which we've already discussed. The skewed distribution of images is honestly a very good touch on AWS WAF's end, and that is exactly why chairs are so difficult for the model to identify. On top of that, something else I noticed was that the images are very uncanny-valley looking, they are synthetic images, not real ones. This makes them deliberately un-googleable, and it defeats models like MobileNetV2 because they train on real images, plus it beats reverse image searching. But this is all undone by the finite, reused, enumerable pool. A perceptual hash lookup table solves it outright, and you can fully enumerate the rare classes. For a solver the CNN is almost beside the point, the lookup table is the real break. This is the fatal flaw in their CAPTCHA: they have clever ML-hostile imagery, but it buys nothing if the pool is small.
So how do you actually make this CAPTCHA better defensively? You might think to add more images, but I would not recommend that. Adding more images would actually help the attacker, a bigger finite pool just helps the attacker's ML and adds labeling. I also wouldn't recommend adding "harder" images. The real fix is to never reuse an image. A fresh per-request generation kills the cache, and it defeats pre-training. There's no stable set to learn from, and the model must generalize to novel images, which is extremely difficult.
The honest result for my model was 76%. There may have been a little wriggle room left, but definitely not much from the dataset I had. I refused to take the memorizer's 95%, as it seems more like a shortcut to me, and that wasn't what I wanted to do anyway. I hope you guys learned something from this post, this one was a wild ride for me. Stay tuned for the next post.