There are a few ways of doing this. In the pre-AI era, this could be done using keywords, with high-ranking results containing more of the searched keywords than lower-ranking ones, but this isn’t really how things are done anymore. More modern approaches use LLMs, leveraging their ability to understand context, so if we search for “fast food chips”, we’ll have high-ranking food results rather than computer chips.

With LLMs in mind, my approach to search is:

• Take each post from my “Following” feed
• Use an LLM to embed the post in an embedding space
• Search for something, being vague or specific
• Return the posts that are close to the search text in the concept space

This is the code:

from atproto import Client

import os

from sentence_transformers import SentenceTransformer
from langchain_community.embeddings import SentenceTransformerEmbeddings
from langchain_chroma import Chroma

# Check environment variables for login exist
try:
	username = os.environ['BSKY_USERNAME']
	password = os.environ['BSKY_PASSWORD']
except Exception:
	raise RuntimeError(
		"Credentials not set in appropriate environment variables!"
	)

# Login
client = Client()
profile = client.login(username, password)
print("Logged in as {}".format(profile.display_name))

# Get the target timeline
target_feed = client.app.bsky.feed.get_timeline().feed 

# We should get the latest 50 posts
print("Feed length: {}".format(len(target_feed)))

posts = []
  
for p in target_feed:
	uri = p.post.uri.split('/')[-1]
	# Get the post text
	text = p.post.record.text
	
	try:
		# Add descriptions if the post has an external link
		text += " - " + p.post.embed.external.description
	except:
		pass
	
	posts.append(text)

# Embed the posts using a sentence transformer  
embedding_model = SentenceTransformer(
	'sentence-transformers/all-mpnet-base-v2'
)

# Create the SentenceTransformerEmbeddings wrapper
embedding_function_wrapper = SentenceTransformerEmbeddings(
	client=embedding_model
)

# Initialise Chroma using the from_texts class method
# Chroma embeds the posts and gives us some nice functions for top-k retrieval
vector_db = Chroma.from_texts(
	collection_name="post_embeddings",
	texts=posts,
	embedding=embedding_function_wrapper,
)

# Search!
while True:
	query = input(">>")
	res = vector_db.similarity_search_with_score(query, k=10)
	
	# Print the top-k results
	for r, s in res:
		print(r.page_content)
		print("Score: {:.2f}".format(s))
		print("="*20)

For this, we’re using the AT Protocol SDK to access Bluesky, logging in using a username and password defined in environment variables (export BSKY_USERNAME="some_username" and export BSKY_PASSWORD="some_password" when using Ubuntu):

# Check environment variables for login exist
try:
	username = os.environ['BSKY_USERNAME']
	password = os.environ['BSKY_PASSWORD']
except Exception:
	raise RuntimeError(
		"Credentials not set in appropriate environment variables!"
	)

# Login
client = Client()
profile = client.login(username, password)
print("Logged in as {}".format(profile.display_name))

Followed by the processing of posts in the feed, which is fairly self-explanatory once you get the right AT protocol function calls, which are not always easy.

The other interesting part is the core of the search function, which uses embeddings of the posts and the query. We define the embedding model:

# Embed the posts using a sentence transformer  
embedding_model = SentenceTransformer(
	'sentence-transformers/all-mpnet-base-v2'
)

# Create the SentenceTransformerEmbeddings wrapper
embedding_function_wrapper = SentenceTransformerEmbeddings(
	client=embedding_model
)

This is a specific model built to embed sentences (all-mpnet-base-v2), which handles a lot of the annoying parts you’d need to implement if you were using a standard LLM, such as chunking the post (so context can be retained across all of the text), and the output being a vector embedding (rather than the generation of new text).

The embeddings are then stored using Chroma:

vector_db = Chroma.from_texts(
	collection_name="post_embeddings",
	texts=posts,
	embedding=embedding_function_wrapper,
)

Chroma is a vector store, designed for exactly this kind of work. The LLM generates the vector representation of the post, we store it using chroma, and we can leverage built-in functions to determine similarities between the stored posts and query text:

res = vector_db.similarity_search_with_score(query, k=10)

The documentation is pretty bad and doesn’t really define the similarity metric or the score. You’d usually use something like cosine similarity when comparing vectors, but high-scoring samples seem to return numbers around 1.5-1.6, which don’t look very cosine.

Anyway, it seems to work okay! The posts that surface from a query are relevant, and on multiple occasions, I’ve been able to find a post I’d been thinking about using a somewhat general query (i.e. not using unusual words that would only appear in that post).

An obvious restriction of this code, as it stands, is that it only takes the most recent 50 posts from your feed. This could quite easily be scaled to search through your entire feed history, but it could take some time to process (embedding each post individually and storing them in Chroma), and while you’re able to access feed histories in Bluesky, I’m not sure how far back they go! But the core methodology of this experiment (LLM post encoding followed by search using Chroma) is built to scale. If we added code to save the vector database to disk and added some feed handling to update the vector store with new posts, we’d have a pretty cool Bluesky search function!

This work started as a curiosity about how AI/ML models perform classification, specifically, why they give the answers they do. As humans, we tend to anthropomorphise AI, which makes it surprising when we realise that it has no clue what it’s actually doing. An image classifier lacks semantic understanding of the images it’s classifying, which can lead to some really strange results (see below). By using explainability methods, we can see why a model assigns a class to an image. While most explanations make sense, they also help us create artificial inputs that have no semantic meaning to humans but still lead to highly confident responses from a model (see below).

Partway through this exploration, I thought, “Hey, allowing people to play with this would be cool!” So, this started the project.

I wanted people to be able to visit a web page, provide input to a classifier, and have the classification results returned to them along with an explanation of those results. From there, people could interpret the explanation and maybe use it to fool the model. In addition to providing a fun little website for people to play with, I’d gain experience deploying an AI/ML model in the real world, along with everything needed to make that a reality. What follows is an explanation of how this project works, the design decisions I made along the way, and what I learned in the process!

Project Outline

There are a few parts to this project. We need: 1) An MNIST classifier with explainability. 2) A way to interface with the MNIST classifier from a web page. 3) A way to publish the app so anyone can access it.

The MNIST classifier isn’t much of a problem. This will just be a small convolutional network built with PyTorch. We’ll train the model, store the weights, and then load it at runtime for inference. Since it’s just a small MNIST model, and the input will be a single $28 \times 28$ grayscale image, we don’t need to involve a GPU at inference time. The inclusion of Integrated Gradients for explainability shouldn’t add too much overhead either, so we should be fine with just CPU compute for this project. Of course, if we were handling large images, lots of inputs, or massive models, a GPU would be a necessity, but I’m being very cost-conscious here and don’t want to hammer in a nail with a sledgehammer. To include explainability with the classifications, I used Captum. A library that simply wraps the PyTorch model, and when given the image we’ve just classified, along with the prediction, it generates an explanation of how the model arrived at its classification. I decided to use the Integrated Gradients¹ approach to model explainability. To compute the integrated gradients of a model $F$ we start with the input to the model $x$ and a baseline $x’$. The baseline is essentially an input that should elicit no meaningful response from the network (i.e., a response as uniform as possible). In the case of images, this would be a black image. We select a dimension $i$ (a pixel in the case of images) and take a path integral of the gradient along a linear interpolation from $x$ to $x’$:

Since $\alpha \in [0,1]$, we gradually move along the path from a completely black image to the original colour version. While this is happening, we’re paying close attention to the rate of change of the model output with respect to $x_i$. This allows us to understand the importance of dimension $i$ to the classification. Doing this for each pixel (or group of pixels) will enable us to see what parts of the input contribute most strongly to the output, producing a heatmap showing how each section contributes to the final class.

Interfacing with the MNIST classifier is an interesting part of this project for me. I’m a classically trained computer scientist, so most of the programs I’ve created in the past have been command-line only. This is often great for efficiency and getting to the core of the problem you’re solving, but it isn’t very accessible to the general public. Because I wanted this to be a web app to exercise my deployment muscle, I knew this would have to be a web page with some JavaScript to make it interactive. I could also use a Python backend to serve the web pages and interact with the model. This was the part I was most cautious about, 1) because I’ve not done much web programming before, and 2) because I’m a security person and know that there are a lot of things that can go wrong security-wise if you’re not careful with your web setup. This is the section I relied on LLMs most heavily. Using them to help me create the JavaScript for the frontend, plug it into the Python backend, and make sure I had addressed most of the obvious security issues. I can’t say the result is perfect, but it does work and I haven’t been hacked (yet, please don’t, if you do find any security issues, please tell me)!

To publish the app, I knew I wanted to use a popular service to gain experience on a well-known platform. I was using Docker for the entire project, so having some way to say: “here’s a Docker container, just make it work” was another aim. I settled on Google Cloud Run, which has the container approach I was looking for, integrated load balancing, good statistics tracking, and is free, provided I don’t exceed a set number of requests. I was a bit disappointed that I didn’t get to do much load balancing because it’s been a while since I’ve done any Kubernetes work, but Google Cloud Run just works, so maybe it’s a good thing not to have another point of failure that I’d have definitely messed up!

MNIST Classifier

Creating the MNIST classifier was fairly straightforward. I used a small architecture I’ve used for a few other projects, consisting of two convolutional layers and one fully connected layer before the output. Since MNIST is such a small dataset, I encountered overfitting issues, so I added some aggressive dropout after the convolution and fully connected layers with a probability of 0.5. The architecture could achieve good results on MNIST (>99%). After training the model weights were saved, and to reduce the project’s attack surface, I removed all training code, leaving only the minimal code to perform inference. The model is loaded and configured to evaluate inputs when the web app starts, so we only load it once per user. Explainability (provided by Captum) is implemented in the backend Python code after classification, with the Captum model initialised after the classifier is loaded from disk.

I admit that the architecture for the classifier is a bit more sophisticated than it needs to be, but I just grabbed a model that I’d previously used for CIFAR classification and wanted to make sure I didn’t run into any capacity issues.

The code for the classifier can be found here.

Real inputs are never simple

For this section I’ll skip forward a little, to after I implemented the interface. Once I had the web app code in place to provide input to the model, and a strong model trained (10 epochs, learning rate of 1e-4, final validation accuracy of 98.87%, loss of 0.0403, and 99.14% on the test set), I could finally test it end-to-end. This, unfortunately, didn’t go well. While the model performed as expected on some inputs, on others (as shown above), it failed to classify them correctly in highly perplexing ways. The integrated gradient explainability provides clues. The first image of the 8 (or 9 according to the model) places a negative weight on the bottom left curve of the input, meaning that it is not important to the classification process at all. A similar thing is happening in the bottom image: the left horizontal cross through the centre of the 7 and the line below it are being ignored. This does help to explain why the model is returning these classes (if we remove these sections of the input, we could argue that the 8 would look like a 9 and the 7 would be closest to a 3), but the question is: why is this happening?

One argument I’ll address quickly is that the number of epochs is quite small. So let’s train longer (even though we achieve > 99% accuracy in just 10 epochs). Increasing the number of epochs to 25 (more epochs started to hint at overfitting), I achieved a validation set accuracy of 99.18%, with a loss of 0.0416, and a test set accuracy of 99.19%. Still not great! Although we do see greater confidence in attributing parts of the input to the classification.

I thought the architecture might be too complex for such a simple task, so I removed the convolutional layers in favour of a simple fully connected network. This didn’t work out well. Even though this network is consistently the most confident, and the explanation is very clear where the classification comes from, it’s completely wrong!

What if we revert to the original architecture and add dropout layers to regularise the model? The new model has a validation accuracy of 98.63%, a loss of 0.0498, and a test set accuracy of 98.73%. This seems to work quite well. But we also have some failures. The major departures in classification we see appear to come from inputs that don’t use the full space. For example, the first 7 in the images above. The MNIST dataset is a perfect encapsulation of the issue we run into when training models to perform a task and then send them out into the wild. That is: people provide weird inputs. 7s are tiny, and on the right-hand side of the input space, 3s or 6s are squashed. We need to make the model more robust to the many ways people can squash their numbers into the input space, and the standard MNIST dataset doesn’t do this. The images are perfectly centred, scaled correctly and all straight. To allow the model to generalise to a wider range of inputs, we need to make the training dataset more varied.

Torchvision provides a really quick way to introduce the changes we want to make to the dataset. We use the RandomAffine transform:

torchvision.transforms.RandomAffine(degrees=30, translate=(0.4, 0.4), scale=(0.3, 1.7))

This function allows us to randomise the application of a variety of affine transforms, such as: Rotation (specified with the degrees parameter), which rotates the input by up to 30 degrees; Translation (translate parameter), which translates the input by a fraction of the size of the input horizontally or vertically (in this example, we translate the image by up to 40% of it’s size up, down, left, or right); Scale (scale parameter), which scales the image between 30-170% of its original size.

These transformations occur during dataset loading. This would be more effective if we applied them when passing the input to the model, but the results seem to work well enough without that extra layer of randomness in the training dataset!

After training for 75 epochs, we have a validation accuracy of 80.05%, with a loss of 0.6576, and a testing set accuracy of 80.27%. While we do see a significant drop in the model’s accuracy, it is more robust to unexpected representations of valid inputs. And on the more general, centred inputs, we were struggling with before: Seems to work pretty well! This experience reminds me that you can never be too set in your ways when implementing any computer program or ML model that interfaces with a user. There’s always going to be something surprising that you never really considered. In this case, it was that the standard MNIST dataset is not very representative of the inputs you should expect from a user. Especially if you give them a lot of space to write in!

UPDATE: I relaxed the transformations (reducing rotation to 15 degrees, and scale to (0.8, 1.2)), and reduced the model size (using fewer convolution kernels and fully connected units). This new model achieved a validation accuracy of ~93%, while reducing the size of the model on disk from 12.5 MB to 1.6 MB!

Building an interface

As I mentioned in the project overview, I didn’t have much experience creating frontends for web apps. I hadn’t really done much work with JavaScript, especially with interacting with users, sending information to a backend, and retrieving something to show on the page afterwards. Because of this, I relied on LLMs to help me get the code right.

I started by defining what the web app should do at a high level. I wanted a small black canvas that visitors could draw on with white ink (reminiscent of the MNIST dataset images). The user should then click a “Classify” button, and the MNIST model should classify the canvas content, returning the class probabilities along with an explanation heat map.

The drawing canvas wasn’t too difficult to get working because people have been drawing on web pages since the beginning of the internet, so a lot of functions already exist to make this easy. Something I hadn’t considered beforehand was adding touch drawing to make the app mobile-friendly. Adding touch drawing just involved a second drawing function, which wasn’t too much extra work. NOTE: The touch implementation is a bit buggy. I really need to look at that again. 😳 The model expects an input image of $28 \times 28$ pixels, but this is a bit small to draw on! To keep things easy to use and process, the drawing canvas is set to $280 \times 280$, and the image is scaled down by a factor of 10 before classification.

Sending the canvas content from the frontend to the backend is done through a POST request to the /classify endpoint in the FastAPI Python code. We pause (await) until we receive a JSON response that includes the classification labels, the softmax probabilities for each label, and the explanation heat map. This data is displayed using Plotly, clearly showing the input, prediction, and explanation for the prediction. All of this data is added to the top of the result stack, pushing all other previous classifications further down the page (I was probably more proud of making that stack work than I should have been!).

The code powering the front end can be found in the GitHub project, specifically here for the JavaScript and here for the HTML. I do not claim that it’s great JavaScript or HTML, but it works.

The backend, acting as a counterpart to the wonderful frontend, was written in Python using the FastAPI web framework. There’s nothing too interesting here. On startup, we load the classifier and create the integrated gradients model, both of which persist throughout the application’s lifetime. Classification happens at the /classify endpoint that takes the image as base64 data, converting it to an image using Pillow. As mentioned earlier, the classification is performed using the previously loaded classifier, with the softmax probabilities, labels, and the explanation heatmap image returned to the frontend.

An interesting part of the backend that I knew I needed to handle, but didn’t know how to approach, was the security of user input to the classifier. Any instance where user input interfaces with code is a security risk, so I wanted to make sure it was handled correctly. So, I asked Claude, which gave this response (imports included for clarity):

from pydantic import BaseModel, Field, field_validator
import base64
import io
from PIL import Image

class ImageRequest(BaseModel):
    image: str = Field(max_length=MAX_PAYLOAD_SIZE)

    @field_validator("image")
    @classmethod
    def validate_image(cls, input_data):
        # Check that it's a PNG data URL (reject anything that isn't)
        if not input_data.startswith("data:image/png;base64,"):
            raise ValueError("User input must be a PNG data URL")

        # First check passed! Next get the base64 data
        _, encoded = input_data.split(",", 1)

        # Validate that base64 is well formed
        try:
            binary_data = base64.b64decode(encoded, validate=True)
        except Exception:
            raise ValueError("Invalid base64 encoding")
        
        # Validate it's actually a PNG by checking for specific bytes
        if not binary_data.startswith(b'\x89PNG\r\n\x1a\n'):
            raise ValueError("Data is not a valid PNG")
        
        # Validate dimensions
        try:
            image = Image.open(io.BytesIO(binary_data))
            w, h = image.size
            if w != MAX_IMAGE_DIM or h != MAX_IMAGE_DIM:
                raise ValueError("Image dimensions exceed maximum size {}x{} > {}x{}".format(w, h, MAX_IMAGE_DIM, MAX_IMAGE_DIM))
            if w < 1 or h < 1:
                raise ValueError("Image too small")
        except ValueError:
            raise
        except Exception:
            raise ValueError("Could not parse image")

        return input_data

The star of the show here is Pydantic, which is a data validation library for Python. We define an ImageRequest class that inherits from BaseModel and has a single field: image. This starts as a string (it’s just a string of bytes), but we make sure that its size does not exceed MAX_PAYLOAD_SIZE, so we’re not ingesting an input that is larger than we expect, which could grind our poor cloud server to a halt (and quickly exceed our maximum memory capacity).

The validation of the input (the validate_image function, which has a decorator @field_validator("image") that I believe connects it to the class field) starts by ensuring the input data is labelled as a PNG. If this is true, the base64-encoded data is decoded into binary, with checks to ensure it’s correctly formatted (binary_data = base64.b64decode(encoded, validate=True)). But anyone can label any data as a PNG, so we have a check to ensure the first few bytes are consistent with the PNG file format (see the Magic number information in the summary thing on the right, which is consistent with the check: binary_data.startswith(b'\x89PNG\r\n\x1a\n')). Now that we know the file size is appropriate, it’s valid base64, and it’s actually a PNG, we can decode the bytes into an image, verify the image size matches what we’d expect, and then continue with classification.

The backend can be found in the main.py file here.

Publishing the app

To make publishing simple, portable, testable, updateable, etc. I Dockerised the program. I’ve been fairly comfortable with Docker for a while, so I knew it was the standard way to build apps that run on the internet. I did extend my knowledge with the build stages, though. As Claude pointed out, leaving all the installs I needed to install the Python libraries in the deployed container would create a larger attack surface than necessary. By building the libraries into a wheel file and passing it from the build to the runtime container, I could make the deployed container more secure and smaller. I also followed best practices to limit user permissions within the container (just in case someone broke out of the program), so that a potential adversary would not have superuser privileges.

The main application was served from the Docker container using Gunicorn, which has lots of options for the number of workers, timeouts, etc. This was bound to a port allowing the app to be served via the same exposed container port. The Dockerfile is here.

To deploy the Docker container, I used Google Cloud Run. This service had a number of attractive selling points, including: autoscaling - so the number of containers that are deployed can grow or shrink depending on demand (there won’t be much), it can automatically build containers from GitHub - so I can have continuous integration (CI) and continuous delivery/deployment (CD), and it handles the TLS/SSL certificate - so my app can use HTTPS. In addition, Cloud Run has access to GPU compute and large-scale storage (external volume mounting), should I ever need to expand the app with GPU-accelerated computing or massive datasets/models in the future.

I kept the setup fairly conservative because this is just a test app, and I don’t expect it to be a massive service that gets lots of traffic:

• I set the minimum number of instances to 0, so if the app doesn’t need to run, it won’t! This also cuts down on costs because I won’t have a container constantly running just in case someone visits. I’m not too worried about cold starts. It does take a bit of time for the app to load sometimes, but that’s a trade-off I’m willing to make.
• The maximum number of instances is 3. I don’t expect many people to access this at once, or at all, so keeping the maximum number low lets me control costs better and make sure expenses don’t get out of hand.
• I gave each instance 512 MiB of memory, and from the service metrics, container memory utilisation is around 80%, so this was a good call! It’s nice to see there’s not much wastage there.
• Each instance has 1 virtual CPU because it doesn’t really need much more. It’s fairly lightweight and performs well.

For the CI/CD portion of this project, I hooked up my GitHub repo to the cloud service to automate the development and deployment process. Cloud Run allows you to add triggers, so when new data is pushed to the repo the Docker container is automatically built, checked and deployed to the service, which streamlines the workflow massively.

Conclusion

So, I have a web app now! You can check it out at letsfoolai.cloud. You can read a bit about how to fool a classifier and then try it out yourself! I don’t expect this to be a very popular app (it’s not that interesting), but I’m so happy it all came together, and it works!

This has been such a learning experience. From the classifier to the web server to the deployment, I’ve had to force myself out of my programming comfort zone to do something new. It’s been a headache at times, primarily in deployment, because that is something that I have no experience with, and understanding all the steps I had to go through to deploy containers and then bug fix what went wrong with very in-depth service analysis pages (I had issues with log file permissions, errors in container deployment due to too little memory, etc.). It has been tough, but so worthwhile! At least now I have a small idea about how modern applications are made available to the masses, and I will sympathise with the engineers if a well-known service is down. 😆

Update on CI/CD

So, it turns out that when I pushed new updates to the repo, they didn’t translate into new builds being deployed by Cloud Run. The culprit was the cloudbuild.yaml file. It wasn’t set up correctly to deploy automatically after a new build was created, so I had to do a manual deploy by opening the deploy options and selecting the build with the correct commit hash. I think this has been fixed with an updated yaml file:

steps:
# Build image
- name: 'gcr.io/cloud-builders/docker'
  args: ['build', '-t', 'gcr.io/${PROJECT_ID}/letsfool-ai:${COMMIT_SHA}', '.']
# Push image to artefact registry
- name: 'gcr.io/cloud-builders/docker'
  args: ['push', 'gcr.io/${PROJECT_ID}/letsfool-ai:${COMMIT_SHA}']
# Deploy to cloud run
- name: 'gcr.io/google.com/cloudsdktool/cloud-sdk'
  entrypoint: gcloud
  args: [
    'run', 'deploy', 'letsfool-ai',
    '--image', 'gcr.io/${PROJECT_ID}/letsfool-ai:${COMMIT_SHA}',
    '--region', 'europe-west1',
    '--platform', 'managed'
  ]

images:
- 'gcr.io/${PROJECT_ID}/letsfool-ai:${COMMIT_SHA}'
options:
  logging: CLOUD_LOGGING_ONLY

We now have separate build and deploy stages, so when I push to the repo, the new container is automatically built and then deployed via Google Cloud Run.

References

Netflix cares about video quality (of course). This post from 2016 details how they created an encoding quality check that more closely aligns with human perception, which is strategically deployed to assess the quality of videos being served to consumers throughout their vast content delivery network. The authors state that this is to evaluate any changes they make to the streaming method, e.g., encoding algorithms, to ensure customers receive high-quality service, which is something I never really considered. So that’s pretty cool.

All of these checks rest on an accurate method for determining the perceptual quality of a video that can be deployed at scale. While human visual perception is the gold standard (humans can quite easily tell if a video stream is low-quality), it’s not very scalable and is difficult to automate in the Netflix content delivery network. So we need to fall back on automated processes to evaluate video quality. Historically, methods such as Mean Squared Error (MSE), Peak signal-to-noise ratio (PSNR), and Structural Similarity Index Measure (SSIM) have been used, but fail to capture the quality of a video encoding from a human perspective.

The aim of this work is to find a correlation between human perception of video quality and the image statistics we can extract.

The authors create a new dataset of 6-second video clips from a range of content available on Netflix and from publicly available sources (animation, videos with film grain, camera motion, indoor/outdoor environments, close-ups of faces, etc.), with a range of resolutions (up to 1920x1080) and bitrates. They then had human observers view each video clip and rate the encoding impairment on a scale from 1 (very annoying) to 5 (not noticeable). The scores from all observers were combined, creating a Differential Mean Opinion Score (DMOS), normalised in the range 0 to 100 (where 100 is the score of the reference image). The collection of videos and DMOS is the NFLX Video Dataset. The dataset is (unsurprisingly) not publicly available (as far as I’m aware). Another interesting point is that they only used H.264/AVC codec for the tests.

The authors found that the results from PSNR, SSIM, Multiscale FastSSIM and PSNR-HSV showed little correlation with the human DMOS scores. Even more interesting is that different types of content exhibit different relationships between the quantitative metric and the DMOS score, with small degradations in any quantitative metric leading to a noticeable drop in the DMOS score.

To align the DMOS scores with a quantitative metric, the authors propose an ML-based approach, Video Multimethod Assessment Fusion (VMAF), that predicts subjective quality from multiple elementary quality metrics. The rationale is that each elementary quality metric has strengths and weaknesses, based on the types of artefacts and degree of distortion. If we combine a whole bunch of different approaches, we should have something better, right?

The authors fuse the following elementary metrics using SVM regression:

• Visual Information Fidelity (VIF) - Image quality metric based on the idea that quality is complementary to the measure of information fidelity loss. Fidelity loss is measured on four scales (a distortion model consisting of blur, additive noise, and global or local contrast changes, I think), with VMAF including each scale as an elementary metric. These are quantified using spectral methods (I believe).
• Detail Loss Metric (DLM) - Image quality metric based on the idea that we can separately measure the loss of details that impact content visibility. This decouples distortions into additive impairments (redundant visual information that does not exist in the original image, e.g. encoding artefacts) and detail losses (loss of useful visual information). Both of these can be quantified by extracting information from the spectral domain (as I understand it).
• Motion - A measure of the temporal difference between adjacent frames, calculated by finding the average absolute pixel difference for the luminance component.

They then train VMAF on the NFLX Training dataset and report results on the testing dataset. The results look convincing. There’s a strong $y=x$ linear correlation between the predicted VMAF and DMOS scores. Compared to the existing Video Quality Model with Variable Frame Delay (VQM-VFD), it beats it in nearly all tests (and is even stronger when VQM-VFD is used to augment it).

This work is available on GitHub - VMAF Development Kit (VDK 1.0.0).

This piece of work was done fairly early on in my PhD. This was back when I was focusing on Geometric Deep Learning and trying to change the world of malware detection with some novel approaches to application representation. The gist of what I wanted to do was to break an app down into basic-blocks and analyse the geometry of the code graph. My feeling was that there may be some unique geometry associated with malware that we can detect, and due to this being a more fundamental part of the structure of a program, it may be more resilient to concept drift. This didn’t end up working out. I worked on this in years 2 and 3 of my PhD and with about a year left to write my thesis I had a massive training pipeline and an AI model that took about a month to train. The approach wasn’t working and I didn’t know which part was the problem. With bugs to fix and only a handful of opportunities to actually train the AI model, I would have been under a lot of pressure to get something working for my thesis (if it was even possible). I decided to cut my losses and change direction. I’ll be making a series of posts about this failed project. Even though it didn’t work, I’m still really proud of what I produced.

⚠️⚠️⚠️

Abstract

The development of malware detection methods in the Android ecosystem has benefited from a huge amount of research in the past few years. Many Machine Learning based approaches have reported incredible performance, but are challenged by real world deployment for two primary reasons. The first is that they are susceptible to dataset drift, occuring as malware evolves to avoid detection which reduces the model performance as a function of time. The second is the lack of transparency in the classifications which inhibits their use in safety and security critical applications.

In this paper we attempt to tackle these two concerns using Graph Neural Network based classifiers and Explainable AI techniques. While there have been many papers proposing Graph Neural Networks as malware classifiers, none have explored the robustness of them to dataset drift, a crucial component to determining their ability to operate effectively in the real world. To explore this point we analyse the ability of Graph Neural Network based classifiers to retain high classification values in the presence of dataset drift. We also explore the use of Explainable AI in analysing the salient graph features which contribute to classification. This provides deep insights into characteristics of malware by providing human interpretable features which increase the verifiability of the classification model and thus the trustworthyness in safety and security critical applications.

Introduction

The Android Operating System (OS) and its many variants are collectively the most popular in the world accounting for approximately 43% of all computer devices¹. This popularity has naturally made it a prime target for malware as adversaries attempt to extract personal information from its many users. Malware typically infects Android devices through the user downloading a malicious application (app) from one of the many app stores available. This has led to a wealth of research focused on classifying apps as benign or malicious based on their properties. Many of these classifiers report high $F_1$, precision and recall values which appears to indicate that this is a solved problem. However, work conducted by Pendlebury et al.² has shown that the reported $F_1$, precision, and recall values often do not tell the full story.

All malware classifiers will suffer from dataset drift. Dataset drift captures many different shifts that can occur in the dataset over time, but it encapsulates the idea that as time progresses the data a classifier is expected to operate on becomes further from the distribution of data it is trained on. Concept drift is of particular interest to us due to the nature of malware, where malicious programs are constantly changing and evolving to avoid detection by classifiers. Testing classifiers in the presence of concept drift is a key evaluation metric which is missing from many of the proposed malware classification methods which presents challenges when it comes to deploying the models in a real world situation.

Approaches to malware classification in the Machine Learning (ML) community have focused on extracting certain properties of the apps which the authors believe would be useful to the categorisation task. Classifiers such as Drebin³ create feature vectors for each app consisting of metadata such as requested hardware components and permissions along with code information such as suspicious API calls and network addresses. Competing ideas in the malware classification space come from Graph Neural Networks (GNNs). Methods such as MAGIC⁴ and MalGraph⁵ argue that we can create more powerful classifiers by representing programs as graphs and use the resulting structure as a feature in addition to features that can be extracted from the code itself. Programs have many natural graph representations, such as Function Call Graphs (FCGs), Control Flow Graphs (CFGs) and Data Flow Graphs (DFGs), therefore this approach to classification has gained traction in the malware detection community. The results reported by many of the GNN malware detection papers show performance competitive with the non-graph variants, however, there is (to our knowledge) no evaluation of GNN methods with respect to concept drift.

We can intuitively reason that GNN approaches to malware classification may be more robust than non-GNN approaches by considering the data we use for classification. Non-GNN classifiers rely on app features which are generally more susceptible to changes in time. For instance both Drebin³ and MaMaDroid⁶ create feature vectors based on API calls (or families of API calls) which can change over time as Android updates the APIs available to developers. On the other hand while GNN methods include API information, this is often augmented by structure information in the form of the program graphs. It is less likely that these graph structures will display significant changes over time which may make them more robust to concept drift.

Explainable Articificial Intelligence (XAI) is a tool that can be used to increase model transparency by highlighting features of the input that are salient to the classification. In the field of malware research, where the safety and security of the system are highly important, this can increase the confidence that end-users have in the model, allowing results to be verified in that we can ensure that the features that are being used for classification are appropriate. In this paper we use an implementation of SubgraphX⁷ to inspect the features that are being used for malware classification which are (hopefully) predominantly subgraph structures. This (hopefully) supports our hypothesis that subgraph structures can be used as robust and reliable features which are characteristic of malware.

Our contributions: This paper evaluates the effect that concept drift has on GNN based malware classifiers in the Android ecosystem. We (hopefully) show that GNN classifiers are more robust to the effects of concept drift, retaining high scores for many months that succeed the training data. We use SubgraphX⁷ to confirm that this behaviour is due to the effective use of the program graph structure which augments the classification process with rich data which is slower to change compared to specific code based features (such as API usage, URL schemes etc.) which are predominantly used by non-GNN based malware classifiers.

Related work

Malware Classification

The automated detection of malware using ML methods has been an area of active research for many years, particularly in the space of Android applications. Popular Android malware classifiers such as Drebin³ and MaMaDroid⁶ apply traditional ML techniques to feature vectors they extract from apps. Drebin focuses on collecting features such as requested hardware access, system permissions, inter/intra-process communications and interfaces to the OS along with restricted API calls and app permissions. This information is then embedded and classified using a Support Vector Machine (SVM). MaMaDroid considers the app program graph in the form of a package or family API call graph extracted from the program binary. The call graph is used to create a Markov chain which models the probability of a transition occurring between API calls. This Markov chain is then used to create a feature vector where entries are the probability of transitioning between the API calls. This is then reduced using Principal Component Analysis (PCA) to an embedding space for classification.

These approaches produced strong $F_1$, precision and recall values, however, they were later shown by Pendlebury et al.² to suffer from concept drift which showed that their effectiveness in real-world applications is limited. This finding could be due to the reliance of these methods on API call analysis which can be volatile between OS updates.

GNNs have been gaining in popularity in the malware classification space, primarily because of their expressive nature of their inputs which allows us to leverage structure of data in the classification process. The most intuitive representation of program data is a program graph which was leveraged by a paper proposing MAGIC⁴. MAGIC focused on classifying benign and malicious programs from Windows binaries using GNNs. The graph representation of the binaries are created by decompiling the programs into assembly, then creating basic-blocks from the instructions. The graph consists of a basic-block as the graph node which contains summary information about the instructions contained within them, with edges representing the connectivity between the basic-blocks in the code structure. This graph is then fed into a Deep Graph Convolutional Neural Network (DGCNN)⁸ with graph reduction layers such as SortPool⁸ before classification which consists of a multi-layer perceptron.

A different approach to the Graph based classification method, which demonstrates the expressive nature of the technology, can be seen in a paper which introduces MalGraph⁵. MalGraph uses a hierarchical graph representation to perform classification using GNNs. At the high level of the hierarchy MalGraph defines an inter-function FCG which details the high level function interactions, and the lower level hierarchy is populated with intra-function CFGs of local functions. This ultimately creates a multi-modal graph where nodes can be local functions which have been processed by a GNN resulting in an embedding to a local function space, or an external function node which is simply one-hot encoded based on the function name. The multi modal FCG is passed to another GNN, followed by a multi-layer perceptron for classification.

One recent application of these graph methods to the Android malware space was presented in a paper by Hei et al. which proposes Hawk⁹. Hawk uses a Heterogeneous Information Network (HIN) which creates a graph of heterogeneous nodes from features such as program APIs, permissions, interfaces etc. with apps being added to the graph with connectivity based on the features they have (e.g. an app node will be connected to a permission node if an app has that permission). A homogeneous graph is then extracted which just contains the app nodes. This is then passed to a variant of the Graph Attention Network (GAT)¹⁰ for classification.

While the graph approaches detailed here all report a strong ability to differentiate malware from goodware on their datasets, there has been (to our knowledge) no thorough analaysis of these methods with respect to concept drift, an issue this paper hopes to change.

Explainable AI

XAI is an invaluable technology in situations where safety and security are of critical importance. In this paper XAI is used to determine features which are salient in the malware classification process. Several papers address the application of XAI to GNNs such as GNNExplainer¹¹, GNN-GI and GNN-LRP¹² and PGExplainer¹³ which all focus on structural node, edge, or node feature explanations of the graphs. Our focus will be on SubgraphX⁷ which specifically uses subgraphs for explanation which, the authors argue, provide a more human interpretable and useful explanation of the classifications.

The application of XAI methods to Graph-based malware classification has been carried out by Herath et al.¹⁴ with CFGExplainer.

$Z = \hat{D}^{-1/2}\hat{A}\hat{D}^{-1/2}XW+b$

Where:

$\hat{A} = A+I$, $\hat{A} \in \mathbb{R}^{N\times N}$ is the adjacency matrix ($N$ is the number of graph vertices/nodes)

$\hat{D}_{ii}=\sum_j^N \hat{A}_{ij}$, $\hat{D} \in \mathbb{R}^{N \times N}$ is the degree matrix

$X \in \mathbb{R}^{N\times F}$ is the feature matrix ($F$ is the number of features we have for each vertex)

$W \in \mathbb{R}^{F\times C}$ is the learned weight matrix ($C$ is the number of weight channels we choose to have)

$b \in \mathbb{R}^{F \times C}$ is the learned bias matrix

Note that the weight and bias part of the above equation is often shortened to $\Theta$ were $\Theta = W+b$.

We’re likely to want to chain multiple graph convolutions together in a ML system, so we re-write the equation above as:

$Z^{(l+1)} = \sigma \left( \hat{D}^{-1/2}\hat{A}\hat{D}^{-1/2}Z^{(l)}\Theta^{(l)} \right)$

Where $l$ denotes the layer and $\sigma$ is a non-linear activation function (usually a sigmoid or ReLU) and $Z^{(0)} = X$.

What we can see here is that each layer of the graph convolution returns another value of $X$, which means we’re operating over the features of the vertices of the graphs. The inclusion of $\hat{D}^{-1/2}\hat{A}\hat{D}^{-1/2}$ shows we’re using some normalised information from the topology of the network to influence this convolution, with this all being weighted by the learned parameters $\Theta$. It’s important to note that $Z$ changes shape as the convolution continues. Initially since $Z^{(0)} = X$, $Z^{(0)} \in \mathbb{R}^{N \times F}$ and later as we multiply by $\Theta \in \mathbb{R}^{F \times C}$ then $Z^{(l)} \in \mathbb{R}^{N \times C}$ which is a matrix of size defined by the number of vertices by the number of weight kernels (each column can be thought of as a way of assigning a value to each node in the graph). $W^{(0)} \in \mathbb{R}^{F \times H}$ (an input to hidden layer with $H$ feature maps) with $W^{(1)} \in \mathbb{R}^{H \times C}$.

Practically, since $\hat{D}^{-1/2}\hat{A}\hat{D}^{-1/2}$ is a constant for graph convolution we can calculate it once and forget about it: $\tilde{A} = \hat{D}^{-1/2}\hat{A}\hat{D}^{-1/2}$

We can then reduce the above equation to a much cleaner form:

$Z^{(l+1)} = \sigma \left( \tilde{A}Z^{(l)}\Theta^{(l)} \right)$

We can easily see how they’ll all chain together:

$Z = \sigma^{(L)}(\dots\sigma^{(1)}(\tilde{A} \sigma^{(0)}(\tilde{A}X\Theta^{(0)})\Theta^{(1)})\dots\Theta^{(L)})$

Looking at $\Theta$

Since $\Theta \in \mathbb{R}^{F \times C}$ then it doesn’t depend on $N$ (the number of vertices). In fact, it doesn’t rely on a specific graph topology at all! This means that we have a graph invariant parameter (which was the focus of much of the early GNN research) which means we can use this learned parameter on many different graphs. This gives credibility to the idea that we can change both the graph and the graph data every training example and still be able to learn a $\Theta$ that can be useful to us!

My first real mind-blowing experience with AI came from DeepMind’s Atari Deep Q-learning agent way back in 2013! Seeing a computer gradually learn to become proficient at a game and then surpass human experts really captured my imagination and was the catalyst for my investment (through education and research) in AI. Since that time, the AI space has exploded with breakthroughs and investment, but with the renewed interest, I feel like the space has lost some of the magic.

In an attempt to find that spark again, I thought I’d embark on a project that aims to be nothing but an interesting exercise for my own amusement. I’ve always been interested in Reinforcement Learning (RL), but haven’t had many opportunities to work in the space (my research is mostly in computer vision/neural network fundamentals), so, to honour DeepMind’s Atari agent, I thought I’d make a computer play a game.

Setup

The game in question is Pokémon Crystal. If you’re unfamiliar with the Pokémon games, you control a character who wanders around a large (mostly linear) world, collecting and battling Pokémon. You collect Pokémon by walking through tall grass in the world, which can trigger random battles with wild Pokémon. In these battles, you can throw Pokéballs which have a percentage chance of capturing the Pokémon you’re battling against (with the percentage dependent on the amount of health the wild Pokémon has, along with multiple other factors). If you successfully capture the Pokémon, you can use it in subsequent battles. Game progress is blocked by ‘trainer battles’, which are battles against other Pokémon trainers. The aim of the game is to beat the 8 gym leaders (more powerful trainers), get the 8 gym badges, then beat the elite 4 (even more powerful trainers), and finally the champion (the most powerful trainer). There are some side stories along the way as well, but that’s the basic premise. So, from an AI perspective, the game can be split into two:

1. Navigation around the world (knowing/finding out where to go next)
2. Beating trainers in Pokémon battles

A few years ago, Peter Whidden published a YouTube video where he trained an RL agent to play Pokémon Red. This is a seriously impressive project that addresses the two halves of the game, with the agent both navigating the world and being competent enough to beat the (not very challenging) CPU-controlled trainers. My attempt takes a much less complete approach. I don’t have the money to pay for either GPUs or GPU time, so I thought I’d work on something smaller. I’ll focus on just winning some battles.

By focusing on a subset of the game (the battles), we can significantly reduce the RL agent’s search space, leading to faster training, lower resource requirements, and faster iteration. We can further reduce the search space by intelligently setting the RL agent’s world state, but I’ll discuss this more later!

If you’re unfamiliar with the Pokémon battle system, it’s a synchronous turn-based system where all users (AI agents, CPU-controlled agents, or human users) select an action at the same time, and those actions then play out over a turn. The actions can include using an attack (that has a limited number of uses), swapping to a different Pokémon, or using an item. Once a turn has played out, the new turn begins (actions are chosen), and this repeats until all of your, or your opponent’s Pokémon have fainted (have 0 health points (HP)). The strategy of Pokémon comes from matching types of Pokémon. Each Pokémon has an elemental type (e.g. Fire, Water, Grass, etc.), with some types being stronger against others. When battling, you want to use a Pokémon that has a type advantage against the opponent to deal maximum damage to them and take less damage yourself. That is the very basic explanation. You can also have different attack types like physical and special, with Pokémon having better defences against physical and special attacks, and then moves that increase the attack/defence ability of your Pokémon, or decrease the attack/defence ability of the opponent, but I’m not skilled enough to write a primer on optimal Pokémon battle techniques (that’s for the agent to find out).

Game emulation

To run the game and pass input to an RL model, I’m using PyBoy, which is a Game Boy/Game Boy Color emulator written in Python. It gives us access to the screen (pixel values in a convenient numpy array), we can send inputs (d-pad, a, b, start, select) and can save and reload states (useful for resetting runs). The emulator can also run at uncapped speeds, allowing us to train rapidly since each battle usually takes a few minutes. All of these features provide us with a suitable environment for automating Pokémon gameplay.

When training an RL agent, we need a metric to assess how well it’s performing and to guide future behaviour with rewards or the promise of a reward. As a first idea, the health of the opponent Pokémon seems like a good place to start. The question is, how do we get this information for the reward?

PyBoy emulates the hardware of the Game Boy Color, recreating the CPU, sound and graphics processor as closely as possible (or as closely as is needed for the program’s aims). This extends to the behaviour and setup of RAM. Since the opponent’s Pokémon’s health is in the game, it should be somewhere in RAM 🙄.

Hacking Pokémon Crystal

The ‘RAM map’ of Pokémon Crystal is quite well documented (I started with datacrystal, then moved on to Glitch City, which seems more complete). From the RAM map, we can find the locations of many pieces of data from the game, e.g. the amount of money the player has (0xd84e-0xd850), or the badges they have from the Johto region (0xd857).

Checking the RAM map, we can easily find the memory address for the opponent Pokémon’s health (0xd216-0xd217).

Now, while we’re here, it’s fun to read program data from RAM. But it’s more fun to write to it! 😏 Tangent time!

Let’s start with something easy. Address 0xdca1 tells the game how many repel steps we have left (repels block wild Pokémon encounters). So, if we change that value in RAM, then we shouldn’t encounter a wild Pokémon! So, I set up the code:

start_time = time.time()
prev_repel = pb.memory[0xdca1]

while pb.tick():
    if time.time() - start_time > 10:
        print("Reset repel")
        pb.memory[0xdca1] = 10
        start_time = time.time()
    
    # Print the repel value if it changes
    if pb.memory[0xdca1] != prev_repel:
        print(pb.memory[0xdca1])
        prev_repel = pb.memory[0xdca1]

This code updates our repel count to 10 every 10 seconds, and prints out the new repel value every time we take a step. What’s interesting is that the game also reports when we run out of repel:

This works!

Let’s do a little more experimentation. Let’s try to force a specific Pokémon (Celebi) to appear during a wild Pokémon battle, and set that Pokémon’s health to 1:

# Set's Pokémon health to 1 across the two bytes that represent it
pb.memory[0xd216] = 0
pb.memory[0xd217] = 1

According to the memory map, 0xd204, 0xd206 and 0xd22e correspond to the ID of the opposing Pokémon. When we start a random battle, those memory values change to the Pokémon’s Pokédex number. So what happens if we overwrite these values with a target Pokédex ID?

pb.memory[0xd204] = 251
pb.memory[0xd206] = 251
pb.memory[0xd22e] = 251

Celebi (Pokédex ID 251) should not be in this patch of grass. In fact, Celebi was a special legendary Pokémon that could only be captured in Ilex Forest (and was only available from a special event in Gold and Silver).

Now we’re comfortable with hacking RAM values, what do we want to achieve? Well, we’d like an agent that can take any team of Pokémon at any level and beat the CPU-controlled opponent in a reasonable battle (e.g., if the agent has level 1 Pokémon and the CPU has level 100, I don’t expect it to win that match-up!).

What do we need to achieve this?

• Have a save state that loads the game just before a battle begins (i.e. a few frames before).
• Set the agent’s and CPU’s party to some arbitrary Pokémon with arbitrary moves and arbitrary levels.
• Provide the RL agent with the world state at the start of each turn.

Let’s get started!

Setting up the AI environment

Battle state

PyBoy allows us to save the state of a game, which captures the contents of RAM, along with the execution environment (program counter, etc.) at any point. I chose to save just before your first battle with your rival. My first attempt was too early: the rival’s Pokémon were set after the first frame was executed, overwriting the hacked team I gave him. After a little trial and error, I managed to save a state just after the rival’s Pokémon were set, which means that we have free rein to set whatever Pokémon we wanted him to have in memory, and it will not be overwritten with any initialisations at the start of the battle.

Setting Agent and CPU Pokémon

This was a little more involved 😅. The teams of Pokémon the agent and CPU can have are nicely arranged in memory, with one Pokémon following on from the previous and all of the properties being consistent in their offset from the first property (Pokémon species (the Pokédex ID)). This makes it easy for us to automate the population of this data.

Apart from the values of start_addr, species_addr and name_addr, which change depending on whether we’re populating the agent or CPU Pokémon and which Pokémon number (1-6), all of the data can be filled in as an offset:

memory[start_addr] = pkm_data['dex_id']
memory[species_addr] = pkm_data['dex_id']

memory[start_addr+1] = 0 # Held item

memory[start_addr+2:start_addr+6] = pkm_data['moves']

# EV values don't matter here
memory[start_addr+11:start_addr+13] = store_integer(0, 2) # HP EV
memory[start_addr+13:start_addr+15] = store_integer(0, 2) # Attk EV
memory[start_addr+15:start_addr+17] = store_integer(0, 2) # Def EV
memory[start_addr+17:start_addr+19] = store_integer(0, 2) # Speed EV
memory[start_addr+19:start_addr+21] = store_integer(0, 2) # Special EV

memory[start_addr+23:start_addr+27] = pkm_data['move_pp']

memory[start_addr+31] = pkm_data['level'] # Pkm level
memory[start_addr+9:start_addr+11] = store_integer(0, 2)

memory[start_addr+33] = 255 # Seems to be related to evolution? Setting this to 255 stops evolution

memory[start_addr+34:start_addr+36] = store_integer(pkm_data['hp'], 2)
memory[start_addr+36:start_addr+38] = store_integer(pkm_data['hp'], 2)
memory[start_addr+38:start_addr+40] = store_integer(pkm_data['attk'], 2)
memory[start_addr+40:start_addr+42] = store_integer(pkm_data['def'], 2)
memory[start_addr+42:start_addr+44] = store_integer(pkm_data['spd'], 2)
memory[start_addr+44:start_addr+46] = store_integer(pkm_data['spatk'], 2)
memory[start_addr+46:start_addr+48] = store_integer(pkm_data['spdef'], 2)

if pkm_data['nickname'] == None:
	pkm_data['nickname'] = "TEST"

nickname = [int(chr_to_hex[s], 16) for s in pkm_data['nickname']]

if len(nickname) > 10:
	long_name = nickname[:10] + [int(0x50)]
else:
	long_name = nickname + [int(0x50)]*(11-len(nickname))

memory[name_addr[1]:name_addr[1]+11] = long_name # 11 chars (10 with terminating hex 50)

In the code above, we’re populating the memory addresses with data from a dictionary (pkm_data), but where does this data come from? As Pokémon fans know, the Pokémon properties (like HP, attack, and defence) depend on the Pokémon itself (some Pokémon have naturally higher attack or defence) and the level. If we want to create teams of Pokémon of arbitrary levels, we need to know how to calculate the correct properties. This is where Pokémon base stats, EVs and IVs come in… 😮‍💨.

Each Pokémon has a series of base stats for HP, attack, defence, special attack, special defence and speed. These base stats kind of tell you what the Pokémon is good at (e.g., a high defence stat tells you that a Pokémon can take a lot of hits, a high speed stat tells you that the Pokémon is likely to make its move first in the turn order). These values change as the Pokémon levels up, but the way they change depends on both the EVs and the IVs. EVs are given to a Pokémon after they defeat another Pokémon in battle, based on the base stats of the Pokémon it defeated (kind of like absorbing the power of enemies you defeat 😬). As the Pokémon defeats enemies, these EVs accrue and contribute to the increase in stats during level up (I don’t think the EVs reset at level up). The maximum EV for each stat is 65535. Pokémon IVs (Individual Values) are kind of like a Pokémon’s genes. They are values specific to each individual Pokémon that determine how well it scales in a particular stat (the maximum IV is 15 for each stat (0b1111)). For example, Umbreon (the best Pokémon) can have stats:

(Data from Serebii).

The variation in stats comes from the EVs and IVs. The equation for calculating a Pokémon’s HP at any level based on base stats, EVs and IVs is:

With all other stats calculated as:

(Equations from Bulbapedia). To keep things simple, when giving the agent or CPU Pokémon, we give them with the maximum stats (EVs of 65535 and IVs of 15). An interesting caveat to note in the calculation is that $\sqrt{\text{EV}}$ is capped at 255.

So, we can see how we can add whatever Pokémon we want to the team, and calculate the correct (maximised) stats at whatever level. But we still have two issues. How do we know the Pokémon’s base stats, and how do we know what moves we can assign to the Pokémon, consistent with its level?

There are many Pokémon databases out there, but the one by veekun does exactly what we need. It just has a bunch of CSV files that we can comb through to get the data we need. In particular, a list of all Pokémon (so we can find the Pokémon name based on the Pokédex ID), a list of Pokémon stats (so we can get the Pokémon base stats), a list of all possible Pokémon moves, and finally a list of moves that each Pokémon can learn, and the levels they learn them at. With some cool SQL inner-joins (that I did have LLM help with as it’s been a while since I’ve done databases), we can query any Pokémon from Gen II (the generation of Crystal), get the Pokémon base stats, and find all possible moves it can learn, and the level it can use those moves at. Here’s how we can create the Pokémon in code:

pkm_info = process_input_pkm([
    ("Umbreon", "Hello", 100, ["tackle", "pursuit", "moonlight", "bite"]),
    ("porygon2", "World", 100, ["conversion-2", "psybeam", "tri-attack", "zap-cannon"]),
    ("skarmory", "This", 100, ["steel-wing", "peck", "fly", "sky-attack"]),
    ("Entei", "Seems", 100, ["flamethrower", "fire-blast", "stomp", "roar"]),
    ("gyarados", "To", 100, ["hydro-pump", "Rain-dance", "Hyper-Beam", "Dragon-Breath"]),
    ("pikachu", "Work", 100, ["Thunderbolt", "dynamic-punch", "thunder", "mega-punch"])
])

This gives us the result:

Getting a world state

Now that we can set the agent and CPU Pokémon arbitrarily, we can re-visit the our todo list:

• Have a save state that loads the game just before a battle begins (i.e. a few frames before).
• Set the agent’s and CPU’s party to some arbitrary Pokémon with arbitrary moves and arbitrary levels.
• Provide the RL agent with the world state at the start of each turn.

Now we just need to provide the RL agent with the world state so it can decide what to do at some regular interval. In Pokémon battles, we have turns, so getting a world state at the start of the turn, then deciding what to do with that turn sounds like an obvious choice. But how do we invoke the agent at the correct time?

This seems like a ridiculous question. There should just be a number somewhere in memory that increments with the turn, right? Well, it doesn’t seem to be as easy as that. First off, there is no documented turn counter reported by datacrystal or Glitch City, and second, as I’ll discuss later, there is no (to my knowledge) clean turn transition point.

To investigate how we can call the agent at the correct time, I created a live memory map using PyGame. The live memory map shows how specified RAM values change over time (my very simple implementation changes the colour of a block when the previous frame had a different value in that RAM address). This should allow us to see which RAM values are associated with the turn transition. The Pan Docs memory map tells us that the Game Boy Color working RAM (WRAM) is between 0xc000 - 0xdfff, which also aligns with the RAM values we’ve been changing to manipulate the game! So I’ll start from 0xc000-0xcfff to keep things manageable:

Now 0xd000-0xdfff:

We can see much more activity in the 0xc... block of memory than 0xd..., but both are quite busy around the turn transition.

After a lot of trial-and-error, narrowing down, and times of pure frustration, I think that memory address 0xc6dc is the turn counter. So, we can grab the world state (Pokémon health information, number of times we can use a move etc.) when the data in location 0xc6dc changes? Of course not! From some experimentation, the turn counter increments before all damage calculations have been done, so the health of the Pokémon is incorrect. This means that we need to wait until the official beginning of the next turn i.e. when the user can actually provide inputs. 😤 Yet another hurdle. Almost like the game was made with little regard to how an AI could be trained on it in the future…

Something that I noticed while inspecting the RAM of the running game, was that when the user selects a different option (Fight, Pkmn, Pack, Run), there is a memory address (0xcfac) that responds to this, storing a different value for each selected value. The fight action is always selected at the start of the turn, clearing whatever was selected at the end of the last turn, meaning that we usually have a transition from some value to 193. So, to detect the start of the turn we can trigger on this transition. The problem is that if we just listen for a transition from any number to 193, we will trigger when we select a different option (e.g. Pkmn) then change our mind and go back to fight. But none of the selections set memory location 0xcfac to 0. So, can we listen for the official turn change measured by 0xc6dc, set the value of 0xcfac to 0, and then listen for the transition from 0 → 193?

Yes. It works… 😵‍💫

Back to the todo list:

• Have a save state that loads the game just before a battle begins (i.e. a few frames before).
• Set the agent’s and CPU’s party to some arbitrary Pokémon with arbitrary moves and arbitrary levels.
• Provide the RL agent with the world state at the start of each turn.

Amazing!

Next steps

You may have been a bit confused by some of the choices being made in that battle. That may be because an agent made those choices at random! So now we need to train it to make intelligent decisions!

This brings you up to date with where I’m currently at. The next steps are going to be the tricky part (as if everything else was simple to do). Most of the future complexity will be around the formal definition of the problem and the architecture of the RL model.

I started this project thinking that I’d define the world state as the observables of the game (agent Pokémon health, amount of uses each of each move of the agent’s Pokémon, opponent health), while this gives us a very restricted search space (which hopefully translates to fast training), it may also be too restricted to allow for generality if we change Pokémon teams. It also means the agent may not be able to learn Pokémon type matchups, since we never explicitly state what our Pokémon is or what the opponent’s is. Time and experiments will tell, but I’m pleased with the progress I’ve been able to make, and look forward to making more mistakes in the future!

This has been a really awesome project so far. I’ve learned a lot about hacking Game Boy Color games (not sure how useful that is), and I really had to think about the design of the pipeline feeding an AI model (in something other than the image-based networks I’m used to), specifically in terms of how we represent the world with data that can be used to find something useful.

I’m still feeling energised to get this project finished and hope to have more to show soon!

This docker container seems to work on my computer (Ubuntu 24 with a 3070 GPU). It will only work on a computer with an NVIDIA GPU and may need the container toolkit.

The container is a Ubuntu 18.04 environment (just like the Wav2Lip one) and installs an old version of CUDA so we can use the older versions of PyTorch etc. Since the container is old, there are some security concerns which (I hope) are addressed by creating a non-priveleged user in the dockerfile who we use to execute all of the code.

I’ve set it all up so when the dockerfile is built and run (code snippets below), you should be left in a bash environment you can use to execute inference.py. The directory ext_files is mounted as an external volume to the container, so you should have access to it (and all files inside it) when in the bash environment. This means that you can put a dataset of images and audio inside ext_files and then pass the file paths to inference.py to generate the videos, which are then saved to the ext_files directory (just remember to change the --outfile file name to avoid overwriting previously generated videos). When you’re finished and exit the container the generated videos should be in the ext_files directory.

Here is the build command:

docker build -t wav2lip -f wav2lip .

And run command:

docker run --rm -v ./ext_files:/home/appuser/work/vol --gpus all --security-opt=no-new-privileges --cap-drop all -it --entrypoint /bin/bash wav2lip

Here’s the code you execute inside the container:

python3 inference.py --checkpoint_path wav2lip_gan.pth --face vol/mona.jpeg --audio vol/ada.mp3 --outfile vol/res.mp4

Here’s the dockerfile:

FROM nvidia/cuda:12.1.0-base-ubuntu18.04

# Create a non-priveleged user
ARG UID=1000
ARG GID=1000
RUN groupadd -g ${GID} appgroup && \
    useradd -m -u ${UID} -g ${GID} -s /bin/bash appuser

# Install necessary packages
RUN apt-get update && apt-get install -y \
    python3-pip python3-dev python3-wheel \
    python3-setuptools python3-apt python3-soundfile \
    build-essential wget \
    git \
    libsm6 libxext6 libxrender1 libglib2.0-0 \
    ffmpeg \
    && rm -rf /var/lib/apt/lists/*

# Make the workspace and give the user access
RUN mkdir -p /home/appuser/work && \ 
	chown -R ${UID}:${GID} /home/appuser

WORKDIR /home/appuser/work

# Clone the Wav2Lip repo
RUN git clone https://github.com/Rudrabha/Wav2Lip /home/appuser/work
# Give the non-priveleged user access
RUN chown -R appuser:appgroup /home/appuser/work

RUN pip3 install --upgrade pip setuptools wheel

# Download the wav2lip model from google drive
RUN pip install gdown
RUN gdown https://drive.google.com/uc?id=1_OvqStxNxLc7bXzlaVG5sz695p-FVfYY

# Download the face detection model
RUN wget -O face_detection/detection/sfd/s3fd.pth https://www.adrianbulat.com/downloads/python-fan/s3fd-619a316812.pth

# Install packages Wav2Lip requires
RUN pip3 install numpy==1.17.1
RUN pip3 install opencv-python==4.1.0.25
RUN pip3 install torch==1.10.1+cu111 torchvision==0.11.2+cu111 torchaudio==0.10.1 -f https://download.pytorch.org/whl/cu111/torch_stable.html
RUN pip3 install tqdm==4.45.0
RUN pip3 install numba==0.43.1
# librosa causes a lot of issues, so don't pin a version and let pip figure out compatibility
RUN pip3 install librosa

USER appuser

Machine Learning and AI models are often viewed as closed-box systems where we provide input and receive an output, and we don’t have a clear idea of how a model arrives at its answer. This lack of transparency can be unnerving depending on the context in which the model is being used, for instance, in safety- and security-critical systems or where major decisions are being made (deciding whether to accept a loan application, etc.). In addition, sometimes it’s just nice to see that a model understands the problem you have assigned to it and is not using some underlying unrelated aspect of the training data to find a shortcut to an answer (leading to failures when deployed in the real world).

Explainable AI methods focus on attributing the prediction of a DNN to input features. In the case of image classification networks, attribution can be assigned to patches of an image (see occlusion attribution) or even down to the pixel level (see integrated gradients).

I have personally used this approach to discuss the effectiveness of models with collaborators, highlighting how datasets are lacking, or questioning the bounds of the task the AI/ML model is required to perform.

Occlusion Attribution

Occlusion-based attribution¹ is the simplest approach to explainable AI for vision models. With this method, we pass an input to the model and record the output. We then occlude a part of the image by adding a black rectangle of a chosen size and then pass the new input to the network, measuring the change between the original output value and the new one. This rectangle is passed over the image (similar to a convolution kernel). By comparing the change in the output values at the different positions of the occluding rectangle, we can determine the importance of parts of the image to the result. The attribution can be represented as a heat map, which leads to a nice visual way of exploring this aspect of models.

Integrated Gradients

Integrated Gradients² were developed by Sundararajan, Taly & Yan in 2017 to address some of the shortcomings of other explainable AI methods, such as Occlusion-based attribution and DeepLift³. The authors posit that two axioms are important when attributing features to a network output. These are:

1. Sensitivity
2. Implementation Invariance

We’ll dive a little more into what these actually mean:

1) Sensitivity - Says that if we pass an input feature and a baseline (a black image in the case of a vision model or the zero embedding vector for text models) to a model, and they produce different outputs, then a non-zero attribution should be assigned to the input feature.

2) Implementation Invariance - The attributions assigned to two functionally equivalent models (models where the input-output pairings are identical, but implementations differ) should be the same.

While I find the sensitivity argument intuitive to understand, I struggle a bit more with the implementation invariance. The authors argue that during backpropagation of the output with respect to the input, the gradients of the model internals get cancelled out, leading to invariance. This makes sense, but it makes me wonder how true this is in practice. The authors state that if an attribution method fails to be implementation invariant, then the attributions may be sensitive to unimportant aspects of the models, e.g. if the model has more degrees of freedom than necessary to represent the function. From my understanding, model capacity (somewhat related to degrees of freedom) is an incredibly difficult metric to capture (and in my experience, models tend to have higher capacity than needed), so to state that implementation invariance relies on this makes me somewhat confused (but as always, I’m quite sure I’ve probably missed a part of the argument!). If models have more degrees of freedom, this means that a change to the output can be attributed to more than one input feature; therefore, two functionally equivalent models can choose to optimise different features to perform their task, leading to functional equivalence, but not implementation invariance (according to the paper).

To compute the Integrated Gradients of a model $F$, the authors start with an input $x$, and baseline $x’$ and select a dimension $i$ (a pixel in the case of images, or a specific feature like number of bedrooms or square footage in the case of tabular information). A path integral of the gradients is then taken along a linear interpolation from $x’$ to $x$. This is summarised as:

The value of $\alpha$ is defined as $\alpha \in [0,1]$, which gradually moves along the path towards the original image $x$.

An added advantage of Integrated Gradients is that unlike Occlusion Attribution, the input images are not changed in a way that can lead to unnatural images (images outside of the distribution the models have been trained to operate on). For Occlusion Attribution, any changes in classification assigned to areas of an image could be due to the out-of-distribution nature of the images instead, causing inflated importance of certain parts of the image.

Attribution in Practice

If we start with a simple MNIST model (two convolution and a fully connected layer, trained to 99% test accuracy), we can see the Occlusion and Integrated Gradients attribution methods in practice.

These are the results for the Occlusion attribution (red indicates positive attribution to the class):

These are the results for the Integrated Gradients attribution:

As we can see these two sets of results are very similar, with the Integrated Gradient approach having slightly more detail in the attribution. We could probably reduce the stride and sliding window shape further for Occlusion attribution to match the detail of the Integrated Gradients, but this starts to defeat the object of occluding blocks of the image (we’d essentially just be turning individual pixels off at that point).

From this we can start to infer some interesting behaviours of the network we’re analysing:

• First, a zero does not need to be a fully enclosed circle (some gaps are allowed), which makes sense since some people do not draw a fully enclosed circle, with the lack of completion occurring at any degree of the circle! It’s interesting to see that this is being taken into account.
• Second, the number 2 (fitting for the second example) is identified if there is a curved horizontal line, with a slanted line emerging from the (approximate) centre of it.
• Finally, the vertical (ish) line of the number 9 has questionable importance during classification. Provided we have a concave curve at the top of the image and some form of mark below it closer to the bottom of the image, it’s quite happy to consider that a 9!

From these findings, we can do some exploitation of the model! See the results below for examples of this.

These images show that we can take the information we discovered about the attribution of features to classes and create examples which have no semantic meaning, but result in high-confidence classifications. The image of the ‘2’ here is particularly striking in just how devoid it is of meaning, yet the model is 100% certain that it’s a 2. This reminds me of rubbish images, discussed in a previous post! From these images, we can find weaknesses in datasets which we may be able to address with more data, correcting the ‘holes’, or we may even discover an edge case in the problem we’re trying to solve. All of which leads to interesting conversations and a more exact definition of the task the AI should perform!

Applying attribution to larger models (Inception V3⁴) pre-trained on the ImageNet dataset, we see similar attribution areas in Bronco the beagle⁵.

Occlusion attribution:

Integrated Gradients:

Both Occlusion and Integrated Gradient attribution highlight the nose, left eye and part of the shoulder as areas that contribute strongly to the classification!

Conclusion

Hopefully this exploration has been interesting! Explainable AI methods provide a really nice visual way to describe what models find important in the classifications they return. From this we can understand where failure modes may exist in underperforming models and address them, either through training on edge cases or a respecification of the task they are being asked to perform.

We’ll focus on the Vision Transformer (ViT)¹, which is a transformer-based image classification architecture that – like the previous architectures I’ve experimented with before – can achieve good classification accuracy on ImageNet and other benchmark datasets. ViT uses a transformer architecture which is as close to the original² as possible, which poses some issues when we move from natural language to images. The original transformer architecture expects inputs to be a finite sequence of one-dimensional tokens, which is challenging if we want to input two-dimensional images. To ameliorate this issue, the authors propose splitting the image into non-overlapping patches, flattening them, passing them through a layer that performs linear projection, taking these patch embeddings and adding a position embedding to add some positional relationship, and finally prepending with the embedding of the class (during training). This converts the image into a sequence of tokens that can then be passed to the transformer architecture (see paper for more details). The interesting part of ViT is that it does not contain any convolution layers, so there are no kernels that we can visualise. We need to find another way to extract learned features!

According to 3b1b³, the interstitial, fully connected layers that occupy the space between the attention layers contain neurons, which represent questions. When we maximise the activation, we find the highest possible response to the question, which should give us some idea about what this question is!

With some minor changes to the visualisation code presented in this post, we can get it working with transformer networks! The most significant adjustments are detailed below:

• Setting opt_type to neuron and changing the extraction process: act = activation[:, filters].
• Changing the image transformations: vit_transform = torchvision.models.ViT_B_16_Weights.IMAGENET1K_V1.transforms(), which are used in the viz_transform composition rather than the googlenet_transform.

If we target a later layer in the transformer (encoder.layers.encoder_layer_11.mlp.3) and optimise the image, we get the following result.

This image doesn’t really give us many answers with respect to the ‘question’ the neuron may be asking. While we do see some patterns emerge, there isn’t really any connection between the random bursts of something vaguely recognisable. In addition, if we optimise a second time, we get this image.

While there are some aspects of both images that are similar, there seems to be a high variance in the patterns being produced. This indicates that the method is not converging to a consistent high activation image.

One reason for this could be the nature of the activation maximisation process itself and how it interacts with the transformer architecture. The activation maximisation process developed in the previous posts works by passing in some random noise, catching the activation of a particular kernel/neuron within the classification network, traversing the gradient in a direction that causes the kernel/neuron to respond more strongly to it, and then backpropagating that change to the input image. However, in the transformer architecture, the ‘context’ of information within the image (extracted by the attention heads) influences the activation of the interstitial ‘question’ neurons. This means that when we optimise the image, not only are we changing the image, but we’re also changing the context. This kind of makes the optimisation process a moving target because the image and the context influence the neuron we’re trying to optimise for, and changing the image changes the context.

This gives me the impression that this is not a suitable way to explore the internal representations of transformers. I’ll need to find an alternative!