1. Introduction

In the previous post, we built an AI-powered chat application on our local computer using microservices. Our stack included FastAPI, Docker, Postgres, Nginx and llama.cpp. The goal of this post is to learn more about the fundamentals of cloud deployment and scaling by deploying our app to Azure, making it available to real users. We’ll use Azure because they offer a free education account, but the process is similar for other platforms like AWS and GCP.

You can check a live demo of the app at chat.jorisbaan.nl. Now, obviously, this demo isn’t very large-scale, because the costs ramp up very quickly. With the tight scaling limits I configured I reckon it can handle about 10–40 concurrent users until I run out of Azure credits. However, I do hope it demonstrates the principles behind a scalable production system. We could easily configure it to scale to many more users with a higher budget.

I give a complete breakdown of our infrastructure and the costs at the end. The codebase is at https://github.com/jsbaan/ai-app-from-scratch.

1.1. Recap: local application

Let’s recap how we built our local app: A user can start or continue a chat with a language model by sending an HTTP request to http://localhost. An Nginx reverse proxy receives and forwards the request to a UI over a private Docker network. The UI stores a session cookie to identify the user, and sends requests to the backend: the language model API that generates text, and the database API that queries the database server.

Table of contents

1. Introduction
   1.1 Recap: local application
2. Cloud architecture
   2.1 Scaling
   2.2 Kubernetes Concepts
   2.3 Azure Container Apps
   2.4 Azure architecture: putting it all together
3. Deployment
   3.1 Setting up
   3.2 PostgreSQL server deployment
   3.3 Azure Container App Environment deployment
   3.4 Azure Container Apps deployment
   3.5 Scaling our Container Apps
   3.6 Custom domain name & HTTPS
4. Resources & costs overview
5. Roadmap
6. Final thoughts
   Acknowledgements
   AI usage

2. Cloud architecture

Conceptually, our cloud architecture will not be too different from our local application: a bunch of containers in a private network with a gateway to the outside world, our users.

However, instead of running containers on our local computer with Docker Compose, we will deploy them to a computing environment that automatically scales across virtual or psychical machines to many concurrent users.

2.1 Scaling

Scaling is a central concept in cloud architectures. It means being able to dynamically handle varying numbers of users (i.e., HTTP requests). Uvicorn, the web server running our UI and database API, can already handle about 40 concurrent requests. It’s even possible to use another web server called Gunicorn as a process manager that employs multiple Uvicorn workers in the same container, further increasing concurrency.

Now, if we want to support even more concurrent request, we could give each container more resources, like CPUs or memory (vertical scaling). However, a more reliable approach is to dynamically create copies (replicas) of a container based on the number of incoming HTTP requests or memory/CPU usage, and distribute the incoming traffic across replicas (horizontal scaling). Each replica container will be assigned an IP address, so we also need to think about networking: how to centrally receive all requests and distribute them over the container replicas.

This “prism” pattern is important: requests arrive centrally in some server (a load balancer) and fan out for parallel processing to multiple other servers (e.g., several identical UI containers).

2.2 Kubernetes Concepts

Kubernetes is the industry standard system for automating deployment, scaling and management of containerized applications. Its core concepts are crucial to understand modern cloud architectures, including ours, so let’s quickly review the basics.

• Node: A physical or virtual machine to run containerized app or manage the cluster.
• Cluster: A set of Nodes managed by Kubernetes.
• Pod: The smallest deployable unit in Kubernetes. Runs one main app container with optional secondary containers that share storage and networking.
• Deployment: An abstraction that manages the desired state of a set of Pod replicas by deploying, scaling and updating them.
• Service: An abstraction that manages a stable entrypoint (the service’s DNS name) to expose a set of Pods by distributing incoming traffic over the various dynamic Pod IP addresses. A Service has multiple types:
  - A ClusterIP Service exposes Pods within the Cluster
  - A LoadBalancer Service exposes Pods to outside the Cluster. It triggers the cloud provider to provision an external public IP and load balancer outside the cluster that can be used to reach the cluster. These external requests are then routed via the Service to individual Pods.
• Ingress: An abstraction that defines more complex rules for a cluster’s entrypoint. It can route traffic to multiple Services; give Services externally-reachable URLs; load balance traffic; and handle secure HTTPS.
• Ingress Controller: Implements the Ingress rules. For example, an Nginx-based controller runs an Nginx server (like in our local app) under the hood that is dynamically configured to route traffic according to Ingress rules. To expose the Ingress Controller itself to the outside world, you can use a LoadBalancer Service. This architecture is often used.

2.3 Azure Container Apps

Armed with these concepts, instead of deploying our app with Kubernetes directly, I wanted to experiment a little by using Azure Container Apps (ACA). This is a serverless platform built on top of Kubernetes that abstracts away some of its complexity.

With a single command, we can create a Container App Environment, which, under the hood, is an invisible Kubernetes Cluster managed by Azure. Within this Environment, we can run a container as a Container App that Azure internally manages as Kubernetes Deployments, Services, and Pods. See article 1 and article 2 for detailed comparisons.

A Container App Environment also auto-creates:

1. An invisible Envoy Ingress Controller that routes requests to internal Apps and handles HTTPS and App auto-scaling based on request volume.
2. An external Public IP address and Azure Load Balancer that routes external traffic to the Ingress Controller that in turn routes it to Apps (sounds similar to a Kubernetes LoadBalancer Service, eh?).
3. An Azure-generated URL for each Container App that is publicly accessible over the internet or internal, based on its Ingress config.

This gives us everything we need to run our containers at scale. The only thing missing is a database. We will use an Azure-managed PostgreSQL server instead of deploying our own container, because it’s easier, more reliable and scalable. Our local Nginx reverse proxy container is also obsolete because ACA automatically deploys an Envoy Ingress Controller.

It’s interesting to note that we literally don’t have to change a single line of code in our local application, we can just treat it as a bunch of containers!

2.4 Azure architecture: putting it all together

Here is a diagram of the full cloud architecture for our chat application that contains all our Azure resources. Let’s take a high level look at how a user request flows through the system.

1. User sends HTTPS request to chat.jorisbaan.nl.
2. A Public DNS server like Google DNS resolves this domain name to an Azure Public IP address.
3. The Azure Load Balancer on this IP address routes the request to the (for us invisible) Envoy Ingress Controller.
4. The Ingress Controller routes the request to UI Container App, who routes it to one of its Replicas where a UI web server is running.
5. The UI web server makes requests to the database API and language model API Apps, who both route it to one of their Replicas.
6. A database API replica queries the PostgreSQL server hostname. The Azure Private DNS Zone resolves the hostname to the PostgreSQL server’s IP address.

3. Deployment

So, how do we actually create all this? Rather than clicking around in the Azure Portal, infrastructure-as-code tools like Terraform are best to create and manage cloud resources. However, for simplicity, I will instead use the Azure CLI to create a bash script that deploys our entire application step by step. You can find the full deployment script including environment variables here 🤖. We will go through it step by step now.

3.1 Setting up

We need an Azure account (I’m using a free education account), a clone of the https://github.com/jsbaan/ai-app-from-scratch repo, Docker to build and push the container images, the downloaded model, and the Azure CLI to start creating cloud resources.

We first create a resource group so our resources are easier to find, manage and delete. The --location parameter refers to the physical datacenter we’ll use to deploy our app’s infrastructure. Ideally, it is close to our users. We then create a private virtual network with 256 IP addresses to isolate, secure and connect our database server and Container Apps.

brew update && brew install azure-cli # for macos

echo "Create resource group"
az group create \
  --name $RESOURCE_GROUP \
  --location "$LOCATION"
  
echo "Create VNET with 256 IP addresses"
az network vnet create \
  --resource-group $RESOURCE_GROUP \
  --name $VNET \
  --address-prefix 10.0.0.0/24 \
  --location $LOCATION

3.2 PostgreSQL server deployment

Depending on the hardware, an Azure-managed PostgreSQL database server costs about $13 to $7000 a month. To communicate with Container Apps, we put the DB server within the same private virtual network but in its own subnet. A subnet is a dedicated range of IP addresses that can have its own security and routing rules.

We create the Azure PostgreSQL Flexible Server with private access. This means only resources within the same virtual network can reach it. Azure automatically creates a Private DNS Zone that manages a hostname for the database that resolves to its IP address. The database API will later use this hostname to connect to the database server.

We will randomly generate the database credentials and store them in a secure place: Azure KeyVault.

echo "Create subnet for DB with 128 IP addresses"
az network vnet subnet create \
  --resource-group $RESOURCE_GROUP \
  --name $DB_SUBNET \
  --vnet-name $VNET \
  --address-prefix 10.0.0.128/25

echo "Create a key vault to securely store and retrieve secrets, \
like the db password"
az keyvault create \
  --name $KEYVAULT \
  --resource-group $RESOURCE_GROUP \
  --location $LOCATION

echo "Give myself access to the key vault so I can store and retrieve \
the db password"
az role assignment create \
  --role "Key Vault Secrets Officer" \
  --assignee $EMAIL \
  --scope "/subscriptions/$SUBSCRIPTION_ID/resourceGroups/$RESOURCE_GROUP/providers/Microsoft.KeyVault/vaults/$KEYVAULT"

echo "Store random db username and password in the key vault"
az keyvault secret set \
  --name postgres-username \
  --vault-name $KEYVAULT  \
  --value $(openssl rand -base64 12 | tr -dc 'a-zA-Z' | head -c 12)
az keyvault secret set \
  --name postgres-password \
  --vault-name $KEYVAULT  \
  --value $(openssl rand -base64 16)

echo "While we're at it, let's already store a secret session key for the UI"
az keyvault secret set \
  --name session-key \
  --vault-name $KEYVAULT  \
  --value $(openssl rand -base64 16)
 
echo "Create PostgreSQL flexible server in our VNET in its own subnet. \
Auto-creates Private DS Zone."
POSTGRES_USERNAME=$(az keyvault secret show --name postgres-username --vault-name $KEYVAULT --query "value" --output tsv)
POSTGRES_PASSWORD=$(az keyvault secret show --name postgres-password --vault-name $KEYVAULT --query "value" --output tsv)
az postgres flexible-server create \
  --resource-group $RESOURCE_GROUP \
  --name $DB_SERVER \
  --vnet $VNET \
  --subnet $DB_SUBNET \
  --location $LOCATION \
  --admin-user $POSTGRES_USERNAME \
  --admin-password $POSTGRES_PASSWORD \
  --sku-name Standard_B1ms \
  --tier Burstable \
  --storage-size 32 \
  --version 16 \
  --yes

3.3 Azure Container App Environment deployment

With the network and database in place, let’s deploy the infrastructure to run containers — the Container App Environment (recall, this is a Kubernetes cluster under the hood).

We create another subnet with 128 IP addresses and delegate its management to the Container App Environment. The subnet should be big enough for every ten new replicas to get a new IP address in the subrange. We can then create the Environment. This is just a single command without much configuration.

echo "Create subnet for ACA with 128 IP addresses."
az network vnet subnet create \
  --resource-group $RESOURCE_GROUP \
  --name $ACA_SUBNET \
  --vnet-name $VNET \
  --address-prefix 10.0.0.0/25

echo "Delegate the subnet to ACA"
az network vnet subnet update \
  --resource-group $RESOURCE_GROUP \
  --vnet-name $VNET \
  --name $ACA_SUBNET \
  --delegations Microsoft.App/environments
  
echo "Obtain the ID of our subnet"
ACA_SUBNET_ID=$(az network vnet subnet show \
  --resource-group $RESOURCE_GROUP \
  --name $ACA_SUBNET \
  --vnet-name $VNET \
  --query id --output tsv)
  
echo "Create Container Apps Environment in our custom subnet.\\
By default, it has a Workload profile with Consumption plan."
az containerapp env create \
  --resource-group $RESOURCE_GROUP \
  --name $ACA_ENVIRONMENT \
  --infrastructure-subnet-resource-id $ACA_SUBNET_ID \
  --location $LOCATION

3.4 Azure Container Apps deployment

Each Container App needs a Docker image to run. Let’s first setup a Container Registry, and then build all our images locally and push them to the registry. Note that we simply copied the model file into the language model image using its Dockerfile, so we don’t need to mount external storage like we did for local deployment in part 1.

echo "Create container registry (ACR)"
az acr create \\
  --resource-group $RESOURCE_GROUP \\
  --name $ACR \\
  --sku Standard \\
  --admin-enabled true
  
echo "Login to ACR and push local images"
az acr login --name $ACR
docker build --tag $ACR.azurecr.io/$DB_API $DB_API
docker push $ACR.azurecr.io/$DB_API
docker build --tag $ACR.azurecr.io/$LM_API $LM_API
docker push $ACR.azurecr.io/$LM_API
docker build --tag $ACR.azurecr.io/$UI $UI
docker push $ACR.azurecr.io/$UI

Now, onto deployment. To create Container Apps we specify their Environment, container registry, image, and the port they will listen to for requests. The ingress parameter regulates whether Container Apps can be reached from the outside world. Our two APIs are internal and therefore completely isolated, with no public URL and no traffic ever routed from the Envoy Ingress Controller. The UI is external and has a public URL, but sends internal HTTP requests over the virtual network to our APIs. We pass these internal hostnames and db credentials as environment variables.

echo "Deploy DB API on Container Apps with the db credentials from the key \
vault as env vars. More secure is to use a managed identity that allows the \
container itself to retrieve them from the key vault. But for simplicity we \
simply fetch it ourselves using the CLI."
POSTGRES_USERNAME=$(az keyvault secret show --name postgres-username --vault-name $KEYVAULT --query "value" --output tsv)
POSTGRES_PASSWORD=$(az keyvault secret show --name postgres-password --vault-name $KEYVAULT --query "value" --output tsv)
az containerapp create --name $DB_API \
  --resource-group $RESOURCE_GROUP \
  --environment $ACA_ENVIRONMENT \
  --registry-server $ACR.azurecr.io \
  --image $ACR.azurecr.io/$DB_API \
  --target-port 80 \
  --ingress internal \
  --env-vars "POSTGRES_HOST=$DB_SERVER.postgres.database.azure.com" "POSTGRES_USERNAME=$POSTGRES_USERNAME" "POSTGRES_PASSWORD=$POSTGRES_PASSWORD" \
  --min-replicas 1 \
  --max-replicas 5 \
  --cpu 0.5 \
  --memory 1

echo "Deploy UI on Container Apps, and retrieve the secret random session \
key the UI uses to encrypt session cookies"
SESSION_KEY=$(az keyvault secret show --name session-key --vault-name $KEYVAULT --query "value" --output tsv)
az containerapp create --name $UI \
  --resource-group $RESOURCE_GROUP \
  --environment $ACA_ENVIRONMENT \
  --registry-server $ACR.azurecr.io \
  --image $ACR.azurecr.io/$UI \
  --target-port 80 \
  --ingress external \
  --env-vars "db_api_url=http://$DB_API" "lm_api_url=http://$LM_API" "session_key=$SESSION_KEY" \
  --min-replicas 1 \
  --max-replicas 5 \
  --cpu 0.5 \
  --memory 1 
  
echo "Deploy LM API on Container Apps"
az containerapp create --name $LM_API \
  --resource-group $RESOURCE_GROUP \
  --environment $ACA_ENVIRONMENT \
  --registry-server $ACR.azurecr.io \
  --image $ACR.azurecr.io/$LM_API \
  --target-port 80 \
  --ingress internal \
  --min-replicas 1 \
  --max-replicas 5 \
  --cpu 2 \
  --memory 4 \
  --scale-rule-name my-http-rule \
  --scale-rule-http-concurrency 2

3.5 Scaling our Container Apps

Let’s take a look at how our Container Apps they scale. Container Apps can scale to zero, which means they have zero replicas and stop running (and stop incurring costs). This is a feature of the serverless paradigm, where infrastructure is provisioned on demand. The invisible Envoy proxy handles scaling based on triggers, like concurrent HTTP requests. Spawning new replicas may take some time, which is called cold-start. We set the minimum number of replicas to 1 to avoid cold starts and the resulting timeout errors for first requests.

The default scaling rule creates a new replica whenever an existing replica receives 10 concurrent HTTP requests. This applies to the UI and the database API. To test whether this scaling rule makes sense, we would have to perform load testing to simulate real user traffic and see what each Container App replica can handle individually. My guess is that they can handle a lot more concurrent request than 10, and we could relax the rule.

3.5.1 Scaling language model inference.

Even with our small, quantized language model, inference requires much more compute than a simple FastAPI app. The inference server handles incoming requests sequentially, and the default Container App resources of 0.5 virtual CPU cores and 1GB memory result in very slow response times: up to 30 seconds for generating 128 tokens with a context window of 1024 (these parameters are defined in the LM API’s Dockerfile).

Increasing vCPU to 2 and memory to 4GB gives much better inference speed, and handles about 10 requests within 30 seconds. I configured the http scaling rule very tightly at 2 concurrent requests, so whenever 2 users chat at the same time, the LM API will scale out.

With 5 maximum replicas, I think this will allow for roughly 10–40 concurrent users, depending on the length of the chat histories. Now, obviously, this isn’t very large-scale, but with a higher budget, we could increase vCPUs, memory and the number of replicas. Ultimately we would need to move to GPU-based inference. More on that later.

3.6 Custom domain name & HTTPS

The automatically generated URL from the UI App looks like https://chat-ui.purplepebble-ac46ada4.germanywestcentral.azurecontainerapps.io/. This isn’t very memorable, so I want to make our app available as subdomain on my website: chat.jorisbaan.nl.

I simply add two DNS records on my domain registrar portal (like GoDaddy). A CNAME record that links my chat subdomain to the UI’s URL, and TXT record to prove ownership of the subdomain to Azure and obtain a TLS certificate.

# Obtain UI URL and verification code
URL=$(az containerapp show -n $UI -g $RESOURCE_GROUP -o tsv --query "properties.configuration.ingress.fqdn")
VERIFICATION_CODE=$(az containerapp show -n $UI -g $RESOURCE_GROUP -o tsv --query "properties.customDomainVerificationId")

# Add a CNAME record with the URL and a TXT record with the verification code to domain registrar
# (Do this manually)

# Add custom domain name to UI App
az containerapp hostname add --hostname chat.jorisbaan.nl -g $RESOURCE_GROUP -n $UI
# Configure managed certificate for HTTPS
az containerapp hostname bind --hostname chat.jorisbaan.nl -g $RESOURCE_GROUP -n $UI --environment $ACA_ENVIRONMENT --validation-method CNAME

Container Apps manages a free TLS certificate for my subdomain as long as the CNAME record points directly to the container’s domain name.

The public URL for the UI changes whenever I tear down and redeploy an Environment. We could use a fancier service like Azure Front Door or Application Gateway to get a stable URL and act as reverse proxy with additional security, global availability, and edge caching.

4. Resources & costs overview

Now that the app is deployed, let’s look at an overview of all the Azure resources it app uses. We created most of them ourselves, but Azure also automatically created a Load balancer, Public IP, Private DNS Zone, Network Watcher and Log Analytics workspace.

Some resources are free, others are free up to a certain time or compute budget, which is part of the reason I chose them. The following resources incur the highest costs:

• Load Balancer (standard Tier): free for 1 month, then $18/month.
• Container Registry (standard Tier): free for 12 months, then $19/month.
• PostgreSQL Flexible Server (Burstable B1MS Compute Tier): free for 12 months, then at least $13/month.
• Container App: Free for 50 CPU hours/month or 2M requests/month, then $10/month for an App with a single replica, 0.5 vCPUs and 1GB memory. The LM API with 2vCPUs, 4GB memory costs about $50 per month for a single replica.

You can see that the costs of this small (but scalable) app can quickly add up to hundreds of dollars per month, even without a GPU server to run a stronger language model! That’s the reason why the app probably won’t be up when you’re reading this.

It also becomes clear that Azure Container Apps is more expensive then I initially thought: it requires a standard-Tier Load balancer for automatic external ingress, HTTPS and auto-scaling. We could get around this by disabling external ingress and deploying a cheaper alternative — like a VM with a custom reverse proxy, or a basic-Tier Load balancer. Still, a standard-tier Kubernetes cluster would have cost at least $150/month, so ACA can be cheaper at small scale.

5. Roadmap

Now, before we wrap up, let’s look at just a few of the many directions to improve this deployment.

Continuous Integration & Continuous Deployment. I would set up a CI/CD pipeline that runs unit and integration tests and redeploys the app upon code changes. It might be triggered by a new git commit or merged pull request. This will also make it easier to see when a service isn’t deployed properly. I would also set up monitoring and alerting to be aware of issues quickly (like a crashing Container App instance).

Lower latency: the language model server. I would load test the whole app — simulating real-world user traffic — with something like Locust or Azure Load Testing. Even without load testing, we have an obvious bottleneck: the LM server. Small and quantized as it is, it can still take up quite a while for lengthy answers, with no concurrency. For more users it would be faster and more efficient to run a GPU inference server with a batching mechanism that collects multiple generation requests in a queue — perhaps with Kafka — and runs batch inference on chunks.

With even more users, we might want several GPU-based LM servers that consume from the same queue. For GPU infrastructure I’d look into Azure Virtual Machines or something more fancy like Azure Machine Learning.

The llama.cpp inference engine is good for single-user CPU-based inference. When moving to a GPU-server, I would look into inference engines more suitable to batch inference, like vLLM or Huggingface TGI. And, obviously, a better (bigger) model for increased response quality — depending on the use case.

6. Final thoughts

I hope this project offers a glimpse of what an AI-powered web app in production may look like. I tried to balance realistic engineering with cutting about every corner to keep it simple, cheap, understandable, and limit my time and compute budget. Sadly, I cannot keep the app live for long since it would quickly cost hundreds of dollars per month. If someone can help with Azure credits to keep the app running, let me know!

Some closing thoughts about using managed services: Although Azure Container Apps abstracts away some of the Kubernetes complexity, it’s still extremely useful to have an understanding of the lower-level Kubernetes concepts. The automatically created invisible infrastructure like Public IPs, Load balancers and ingress controllers add unforeseen costs and make it difficult to understand what’s going on. Also, ACA documentation is limited compared to Kubernetes. However, if you know what you’re doing, you can set something up very quickly.

Acknowledgements

I heavily relied on the Azure docs, and the ACA docs in particular. Thanks to Dennis Ulmer for proofreading and Lucas de Haas for useful discussion.

AI usage

I experimented a bit more with AI tools compared to part 1. I used Pycharm’s CoPilot plugin for code completion and had quite some back-and-forth with ChatGPT to learn about the Azure or Kubernetes ecosystem, and to spar about bugs. I double-checked everything in the docs and most of the information was solid. Like part 1, I did not use AI to write this post, though I did use ChatGPT to paraphrase some bad-running sentences.

Designing, Building & Deploying an AI Chat App from Scratch (Part 2) was originally published in TDS Archive on Medium, where people are continuing the conversation by highlighting and responding to this story.

Introduction

The aim of this project is to learn about the fundamentals of modern, scalable web applications by designing, building and deploying an AI-powered chat app from scratch. We won’t use fancy frameworks or commercial platforms like ChatGPT. This will provide a better understanding of how real-world systems may work under the hood, and give us full control over the language model, infrastructure, data and costs. The focus will be on engineering, backend and cloud deployment, rather than the language model or a fancy frontend.

This is part 1. We will design and build a cloud-native app with several APIs, a database, private network, reverse proxy, and simple user interface with sessions. Everything runs on our local computer. In part 2, we will deploy our application to a cloud platform like AWS, GCP or Azure with a focus on scalability so actual users can reach it over the internet.

You can find the codebase at https://github.com/jsbaan/ai-app-from-scratch. Throughout this post I will link to specific lines of code with this hyperlink robot 🤖 (try it!)

Microservices and APIs

Modern web applications are often built using microservices — small, independent software components with a specific role. Each service runs in its own Docker container — an isolated environment independent of the underlying operating system and hardware. Services communicate with each other over a network using REST APIs.

You can think of a REST API as the interface that defines how to interact with a service by defining endpoints — specific URLs that represent the possible resources or actions, formatted like http://hostname:port/endpoint-name Endpoints, also called paths or routes, are accessed with HTTP requests that can have various types like GET to retrieve data or POST to create data. Parameters can be passed in the URL itself or in the request body or header.

Architecture

Let’s make this more concrete. We want a web page where users can chat with a language model and come back to their previous chats. Our architecture will look like this:

The above architecture diagram shows how a user’s HTTP request to localhost on the left flows through the system. We will discuss and set up each individual service, starting with the backend services on the right. Finally, we discuss communication, networking and container orchestration.

The structure of this post follows the components in our architecture (click to jump to the section):

1. Language model API. A llama.cpp language model inference server running the quantized Qwen2.5–0.5B-Instruct model 🤖.
2. PostgreSQL database server. A database that stores chats and messages 🤖.
3. Database API. A FastAPI and Uvicorn Python server that queries the PostgreSQL database 🤖.
4. User interface. A FastAPI and Uvicorn Python server that serves HTML and support session-based authentication 🤖.
5. Private Docker network. For communication between microservices 🤖.
6. Nginx reverse proxy. A gateway between the outside world and network-isolated services 🤖.
7. Docker Compose. A container orchestration tool to easily run manage our services together 🤖.

1. Language model API

Setting up the actual language model is pretty easy, nicely demonstrating that ML engineering is usually more about engineering than ML. Since I want our app to run on a laptop, model inference should be fast and CPU-based with low memory.

I looked at several inference engines, like Fastchat with vLLM or Huggingface TGI, but went with llama.cpp because it’s popular, fast, lightweight and supports CPU-based inference. Llama.cpp is written in C/C++ and conveniently provides a Docker image with its inference engine and a simple web server that implements the popular OpenAI API specification. It comes with a basic UI for experimenting, but we’ll build our own UI shortly.

As for the actual language model, I chose the quantized Qwen2.5–0.5B-Instruct model from Alibaba Cloud, whose responses are surprisingly coherent given how small it is.

1.1 Running the language model API

The beauty of containerized applications is that, given a Docker image, we can have it running in seconds without installing any packages. The docker run command below pulls the llama.cpp server image, mounts the model file that we downloaded earlier to the container’s filesystem, and runs a container with the llama.cpp server listening for HTTP requests at port 80. It uses flash attention and has max generation length of 512 tokens.

docker run
	--name lm-api \
	--volume $PROJECT_PATH/lm-api/gguf_models:/models \
	--publish 8000:80 \ # add this to make the API accessible on localhost
	ghcr.io/ggerganov/llama.cpp:server 
		-m /models/qwen2-0_5b-instruct-q5_k_m.gguf --port 80 --host 0.0.0.0 --predict 512 --flash-attn

Ultimately we will use Docker Compose to run this container together with the others 🤖.

1.2 Accessing the container

Since Docker containers are completely isolated from everything else on their host machine, i.e., our computer, we can’t reach our language model API yet.

However, we can break through a bit of networking isolation by publishing the container’s port 80 to our host machine’s port 8000 with --publish 8000:80 in the docker run command. This makes the llama.cpp server available at http://localhost:8000.

The hostname localhost resolves to the loopback IP address 127.0.0.1 and is part of the loopback network interface that allows a computer to communicate with itself. When we visit http://localhost:8000, our browser sends an HTTP GET request to our own computer on port 8000, which gets forwarded to the llama.cpp container listening at port 80.

1.3 Testing the language model API

Let’s test the language model server by sending a POST request with a short chat history.

curl -X POST <http://localhost:8000/v1/chat/completions> \
-H "Content-Type: application/json" \
-d '{
	"messages": [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "assistant", "content": "Hello, how can I assist you today?"},
    {"role": "user", "content": "Hi, what is an API?"}
	],
	"max_tokens": 10
	}'

The response is JSON and the generated text is under choices.message.content: “An API (Application Programming Interface) is a specification…”.

Perfect! Ultimately, our UI service will be the one to send requests to the language model API, and define the system prompt and opening message 🤖.

2. PostgreSQL database server

Next, let’s look into storing chats and messages. PostgreSQL is a powerful, open-source relational database and running a PostgreSQL server locally is just another another docker run command using its official image. We’ll pass some extra environment variables to configure the database name, username, and password.

docker run --name db --publish 5432:5432 --env POSTGRES_USER=myuser --env POSTGRES_PASSWORD=mypassword postgres

After publishing port 5432, the database server is available on localhost:5432. PostgreSQL uses its own protocol for communication and doesn’t understand HTTP requests. We can use a database client like psql to test the connection.

pg_isready -U joris -h localhost -d postgres
> localhost:5432 - accepting connections

When we deploy our application in part 2, we will use a database managed by a cloud provider to make our lives easier and add more security, reliability and scalability. However, setting one up locally like this is useful for local development and, perhaps later on, integration tests.

3. Database API

Databases often have a separate API server sitting in front to control access, enforce extra security, and provide a simple, standardized interface that abstracts away the database’s complexity.

We will build this API from scratch with FastAPI, a modern framework for building fast, production-ready Python APIs. We will run the API with Uvicorn, a high-performance Python web server that handles things like network communication and simultaneous requests.

3.1 Quick FastAPI example

Let’s quickly get a feeling for FastAPI and look at a minimal example app with a single GET endpoint /hello.

from fastapi import FastAPI

# FastAPI app object that the Uvicorn web server will load and serve
my_app = FastAPI()

# Decorator telling FastAPI that function below handles GET requests to /hello
@my_app.get("/hello") 
def read_hello():
  # Define this endpoint's response
    return {"Hello": "World"}

We can serve our app at http://localhost:8080 by running the Uvicorn server.

uvicorn main.py:my_app --host 0.0.0.0 --port 8080

If we now send a GET request to our endpoint by visiting http://localhost:8080/hello in our browser, we receive the JSON response {"Hello": "World"} !

3.2 Connecting to the database

On to the actual database API. We define four endpoints in main.py 🤖 for creating or fetching chats and messages. You get a nice visual summary of these in the auto-generated docs, see below. The UI will call these endpoints to process user data.

The first thing we need to do is to connect the database API to the database server. We use SQLAlchemy, a popular Python SQL toolkit and Object-Relational Mapper (ORM) that abstracts away writing manual SQL queries.

We establish this connection in database.py 🤖 by creating the SQLAlchemy engine with a connection URL that includes the database hostname, username and password (remember, we configured these by passing them as environment variables to the PostgreSQL server). We also create a session factory that creates a new database session for each request to the database API.

3.3 Interacting with the database

Now let’s design our database. We define two SQLAlchemy data models in models.py 🤖 that will be mapped to actual database tables. The first is a Message model 🤖 with an id, content, speaker role, owner_id, and session_id (more on this later). The second is a Chat model 🤖, which I’ll show here to get a better feeling for SQLAlchemy models:

class Chat(Base):
  __tablename__ = "chats"

  # Unique identifier for each chat that will be generated automatically.
  id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)

  # Username associated with the chat. Index is created for faster lookups.
  username = Column(String, index=True)

  # Session ID associated with the chat. 
  # Used to "scope" chats, i.e., users can only access chats from their session.
  session_id = Column(String, index=True)

  # The relationship function links the Chat model to the Message model.
  # The back_populates flag creates a bidirectional relationship.
  messages = relationship("Message", back_populates="owner")

Database tables are typically created using migration tools like Alembic, but we’ll simply ask SQLAlchemy to create them in main.py 🤖.

Next, we define the CRUD (Create, Read, Update, Delete) methods in crud.py 🤖. These methods use a fresh database session from our factory to query the database and create new rows in our tables. The endpoints in main.py will import and use these CRUD methods 🤖.

3.4 Endpoint request and response validation

FastAPI is heavily based on Python’s type annotations and the data validation library Pydantic. For each endpoint, we can define a request and response schema that defines the input/output format we expect. Each request to or response from an endpoint is automatically validated and converted to the right data type and included in our API’s automatically generated documentation. If something about a request or response is missing or wrong, an informative error is thrown.

We define the Pydantic schemas for the database-api in schemas.py 🤖 and use them in the endpoint definitions in main.py 🤖, too. For example, this is the endpoint to create a new chat:

@app.post("/chats", response_model=schemas.Chat)
async def create_chat(chat: schemas.ChatCreate, db: Session = Depends(get_db)):
    db_chat = crud.create_chat(db, chat)
    return db_chat

We can see that it expects a ChatCreate request body and Chat response body. FastAPI verifies and converts the request and response bodies according to these schemas 🤖:

class ChatCreate(BaseModel):
    username: str
    messages: List[MessageCreate] = []
    session_id: str

class Chat(ChatCreate):
    id: UUID
    messages: List[Message] = []

Note: our SQLAlchemy models for the database should not be confused with these Pydantic schemas for endpoint input/output validation.

3.5 Running the database API

We can serve the database API using Uvicorn, making it available at http://localhost:8001.

cd $PROJECT_PATH/db-api
uvicorn app.main.py:app --host 0.0.0.0 --port 8001

To run the Uvicorn server in its own Docker container, we create a Dockerfile 🤖 that specifies how to incrementally build the Docker image. We can then build the image and run the container, again making the database API available at http://localhost:8001 after publishing the container’s port 80 to host port 8001. We pass the database credentials and hostname as environment variables.

docker build --tag db-api-image $PROJECT_PATH/db-api
docker run --name db-api --publish 8001:80 --env POSTGRES_USERNAME=<username> --env POSTGRES_PASSWORD=<password> --env POSTGRES_HOST=<hostname> db-api-image

4. User interface

With the backend in place, let’s build the frontend. A web interface typically consists of HTML for structure, CSS for styling and Javascript for interactivity. Frameworks like React, Vue, and Angular use higher-level abstractions like JSX that can ultimately be transformed into HTML, CSS, and JS files to be bundled and served by a web server like Nginx.

Since I want to focus on the backend, I hacked together a simple UI with FastAPI. Instead of JSON responses, its endpoints now return HTML based on template files that are rendered by Jinja, a templating engine that replaces variables in the template with real data like chat messages.

To handle user input and interact with the backend (e.g., retrieve chat history from the database API or generate a reply via the language model API), I’ve avoided JavaScript altogether by using HTML forms 🤖 that trigger internal POST endpoints. These endpoints then simply use Python’s httpx library to make HTTP requests 🤖.

Endpoints are defined in main.py 🤖, HTML templates are in the app/templates directory 🤖, and the static CSS file for styling the pages is in the app/static directory 🤖. FastAPI serves the CSS file at http://hostname/static/style.css so the browser can find it.

4.1 Homepage

The homepage allows users to enter their name to start or return to a chat 🤖. The submit button triggers a POST request to the internal /chats endpoint with username as form parameter, which calls the database API to create a new chat and then redirects to the Chat Page 🤖.

4.2 Chat Page

The chat page calls the database API to retrieve the chat history 🤖. Users can then enter a message that triggers a POST request to the internal /generate/{chat_id} endpoint with the message as form parameter 🤖.

The generate endpoint calls the database API to add the user’s message to the chat history, and then the language model API with the full chat history to generate a reply 🤖. After adding the reply to the chat history, the endpoint redirects to the chat page, that again retrieves and displays the latest chat history. We send POST request to the LM API using httpx, but we could use a more standardized LM API package like langchain to invoke its completion endpoint.

4.3 Authentication & user sessions

So far, all users can access all endpoints and all data. This means anyone can see your chat given your username or chat id. To remedy that, we will use session-based authentication and authorization.

We will store a first party & GDRP compliant signed session cookie in the user’s browser 🤖. This is just an encrypted dict-like object in the request/response header. The user’s browser will send that session cookie with each request to our hostname, so that we can identify and verify a user user and show them their own chats only.

As an extra layer of security, we “scope” the database API such that each chat row and each message row in the database contains a session id. For each request to the database API, we include the current user’s session id in the request header and query the database with both the chat id (or username) AND that session id. This way, the database can only ever return chats for the current user with its unique session id 🤖.

4.4 Running the UI

To run the UI in a Docker container, we follow the same recipe as the database API, adding the hostnames of the database API and language model API as environment variables.

docker build --tag chat-ui-image $PROJECT_PATH/chat-ui
docker run --name chat-ui --publish 8002:80 --env LM_API_URL=<hostname1> --env DB_API_URL=<hostname2> chat-ui-image

How do we know the hostnames of the two APIs? We will look networking and communication next.

5. Private Docker network

Let’s zoom out and take a look at our architecture again. By now, we have four containers: the UI, DB API, LM API, and PostgreSQL database. What’s missing is the network, reverse proxy and container orchestration.

Until now, we used our computer’s localhost loopback network to send requests to an individual container. This was possible because we published their ports to our localhost. However, for containers to communicate with each other, they must be connected to the same network and know each others hostname/IP address and port.

We will create a user-defined bridge Docker network that provides automatic DNS resolution. This means that container names are resolved to their dynamic container’s IP address. The network also provides isolation, and therefore security: you have to be on the same network to be able to reach our containers.

docker network create --driver bridge chat-net

We connect all containers to it by adding --network chat-net to their docker run command. Now, the database API can reach the database at db:5432 and the UI can reach the database API at http://db-api and the language model API at http://lm-api. Port 80 is default for HTTP requests so we can omit it.

6. Nginx reverse proxy

Now, how do we — the user — reach our network-isolated containers? During development we published the UI container port to our localhost, but in a realistic scenario you typically use a reverse proxy. This is another web server that acts like a gateway, forwarding HTTP requests to containers in their private network, enforcing security and isolation.

Nginx is a web server often used as reverse proxy. We can easily run it using its official Docker image. We also mount a configuration file 🤖 in which we specify how Nginx should route incoming requests. As an example, the simplest possible configuration forwards all requests (location / ) from Nginx container’s port 80 to the UI container at http://chat-ui.

http { server { listen 80; location / { proxy_pass http://chat-ui } } }

Since the Nginx container is in same private network, we can’t reach it either. However, we can publish its port so it becomes the only access point of our entire app 🤖. A request to localhost now goes to the Nginx container, who forwards it to the UI and the UI’s response back to us.

docker run --network chat-net --publish 80:80 --volume $PROJECT_PATH/nginx.conf:/etc/nginx/nginx.conf nginx

In part 2 we will see that these gateway servers can also distribute incoming requests over copies of the same containers (load balancing) for scalability; enable secure HTTPS traffic; and do advanced routing and caching. We will use an Azure-managed reverse proxy rather this Nginx container, but I think it’s very useful to understand how they work and how to set one up yourself. It can also be significantly cheaper compared to a managed reverse proxy.

7. Docker Compose

Let’s put everything together. Throughout this post we manually pulled or built each image and ran its container. However, in the codebase I’m actually using Docker Compose: a tool designed to define, run and stop multi-container applications on a single host like our computer.

To use Docker Compose, we simply specify a compose.yml file 🤖 with build and run instructions for each service. A cool feature is that it automatically creates a user-defined bridge network to connect our services. Docker DNS will resolve the service names to container IP addresses.

Inside the project directory we can start all services with a single command:

docker compose up --build

Final thoughts

That wraps it up! We built an AI-powered chat web application that runs on our local computer, learning about microservices, REST APIs, FastAPI, Docker (Compose), reverse proxies, PostgreSQL databases, SQLAlchemy, and llama.cpp. We’ve built it with a cloud-native architecture in mind so we can deploy our app without changing a single line of code.

We will discuss deployment in part 2 and cover Kubernetes, the industry-standard container orchestration tool for large-scale applications across multiple hosts; Azure Container Apps, a serverless platform that abstracts away some of Kubernetes’ complexities; and concepts like load balancing, horizontal scaling, HTTPS, etc.

Roadmap

There is a lot we could do to improve this app. Here are some things I would work on given more time.

Language model. We now use a very general instruction-tuned language model as virtual assistant. I originally started this project to have a “virtual representation of me” on my website for visitors to discuss my research with, based on my scientific publications. For such a use case, an important direction is to improve and tweak the language model output. Perhaps that’ll become a part 3 of this series in the future.

Frontend. Instead of a quick FastAPI UI, I’d build a proper frontend using something like React, Angular or Vue to allow things like steaming LM responses and dynamic views rather than reloading the page every time. A more lightweight alternative that I’d like to experiment with is htmx, a library that provides modern browser features directly from HTML rather than javascript. It would be pretty straightforward to implement LM response streaming, for example.

Reliability. To make the system more mature, I’d add unit and integration tests and a better database setup like Alembic allowing for migrations.

Acknowledgements

Thanks to Dennis Ulmer, Bryan Eikema and David Stap for initial feedback or proofreading.

AI usage

I used Pycharm’s CoPilot plugin for code completion, and ChatGPT for a first version of the HTML and CSS template files. Towards the end, I started experimenting more with debugging and sparring too, which proved surprisingly useful. For example, I used it to learn about Nginx configurations and session cookies in FastAPI. I did not use AI to write this post, though I did use ChatGPT to paraphrase a few bad-running sentences.

Here are some additional resources that I found useful during this project.

• Django vs Flask vs FastAPI YouTube video by Patrick Loeber.
• The FastAPI tutorial was by far most helpful.
• Docker documentation, especially the networking overview, bridge network intro, docker-compose quickstart and networks in docker-compose.
• SQLAlchemy docs and their Object-Relational Mapper (ORM) tutorial.
• PostgreSQL Cheat Sheet.

Designing, Building & Deploying an AI Chat App from Scratch (Part 1) was originally published in TDS Archive on Medium, where people are continuing the conversation by highlighting and responding to this story.

7 surprising things you might not know about the Transformer

Introduction

To get intimately familiar with the nuts and bolts of transformers I decided to implement the original architecture in the “Attention is all you need” paper from scratch. I thought I knew everything there was to know, but to my own surprise, I encountered several unexpected implementation details that made me better understand how everything works under the hood.

The goal of this post is not discuss the entire implementation — there are plenty of great resources for that — but to highlight seven things that I found particularly surprising or insightful, and that you might not know about. I will make this concrete by pointing to specific lines in my code using this hyperlink robot 🤖 (try it!). The code should be easy to understand: it’s well documented and automatically unit tested and type checked using Github Actions.

The structure of this post is simple. The first three points revolve around implementing multi-head attention; the last four are about other components. I assume you have conceptual familiarity with the transformer and multi-head attention (if not; a great starting point is the Illustrated Transformer) and kick off with something that helped me tremendously in better understanding the mechanics behind multi-head attention. It’s also the most involved point.

Table of Contents

Introduction
1. Multi-head attention is implemented with one weight matrix
2. The dimensionality of the key, query and value vectors is not a hyperparameter
3. Scaling in dot-product attention avoids extremely small gradients
4. The source embedding, target embedding AND pre-softmax linear share the same weight matrix
5. During inference we perform a decoder forward pass for every token; during training we perform a single forward pass for the entire sequence
6. Transformers can handle arbitrarily long sequences, in theory
7. Transformers are residual streams
Wrapping up

1. Multi-head attention is implemented with one weight matrix

Before we dive into it; recall that for each attention head we need a query, key and value vector for each input token. We then compute attention scores as the softmax over the scaled dot product between one query and all key vectors in the sentence (🤖). The equation below computes an attention-weighted average over all value vectors for each input token at once. Q is the matrix that stacks query vectors for all input tokens; K and V do the same for key and value vectors.

How, then, do we obtain these query, key and value vectors efficiently for all tokens and heads? It turns out we can do this in one go using a single weight matrix W. This is different from the three projection weight matrices one might expect after reading the paper or the Illustrated Transformer. Let’s walk through how this works.

Let’s say that our input consists of 4 tokens: ["hey", "how", "are", "you"], and our embedding size is 512. Ignoring batches for now, let X be the 4x512 matrix stacking token embeddings as rows.

Let W be a weight matrix with 512 rows and 1536 columns. We will now zoom in on this 512x1536 dimensional weight matrix W (🤖) to find out why we need 1536 dimensions and how multiplying it with X results in a matrix P (for projections) that contains all the query, key and value vectors we need. (In the code 🤖 I call this matrix qkv)

The Matrix Multiplication Behind Multi-head Attention

Each element in the resulting 4x1536 matrix P=X@W is the sum of the element-wise product (in other words: the dot product) between a row vector in X (an embedding) and a column vector in W (some weights).

As a refresher on matrix multiplication, the image below visualizes how to compute the first element of a simple 3x3 matrix when multiplying two 3x3 matrices. The same strategy applies to our bigger matrices X and W.

So, each element in the ith row P[i, :] in our projection matrix is a linear combination of the ith token embedding X[i, :] and one of the weight columns in W. This means we can simply stack more columns in our weight matrix W to create more independent linear combinations (scalars) of each token embedding in X. In other words, each element in P is a different scalar "view" or "summary" of a token embedding in X, weighted by a column in W. This is key in understanding how eight “heads” with “query”, “key” and “value” vectors hide within each of P's rows.

Uncovering attention heads and query, key and value vectors

We can decompose the 1536 columns that we have chosen for W (and end up as the number of columns in P) into 1536 = 8 * 3 * 64. We now uncovered eight heads with each three 64-dimensional vectors hiding in every row in P! Each such “vector” or “chunk” consists of 64 different weighted linear combination of a token embedding and we choose to interpret them in a certain way. You can see a visual representation of P and how to decompose it in the image below. The decomposition also happens in code (🤖).

For multiple sentences in a batch, simply imagine a third dimension “behind” P that turns the 2D matrix into 3D matrix.

Encoder-decoder attention

For encoder-decoder attention this is slightly more involved. Recall that encoder-decoder attention allows each decoder to attend to the embeddings outputted by the topmost encoder

For encoder-decoder attention, we need query vectors for the decoder token embeddings, and key and value vectors for the topmost encoder token embeddings. That’s why we split W into two — a 512x512 and a 512x1024 matrix (🤖) — and perform two separate projections: one to obtain the key and value vectors from the encoder’s embeddings (🤖), and one to obtain query vectors for the decoder’s embeddings (🤖).

Finally, note that we do need a second weight matrix (🤖) in multi-head attention to mix the value vectors from each head and obtain a single contextual embedding per token (🤖).

2. The dimensionality of the key, query and value vectors is not a hyperparameter

I never really thought about this, but I always assumed that the dimensionality of the query, key and value vectors was a hyperparameter. As it turns out, it is dynamically set to the number of embedding dimensions divided by the number of heads: qkv_dim = embed_dim/num_heads = 512/8 = 64(🤖).

This seems like a design choice by Vaswani et al. to keep the number of parameters in multi-head attention constant, regardless of the number of heads one chooses. While you might expect the number of parameters to grow with more heads, what actually happens is that the dimensionality of the query, key and value vectors decreases.

If we look at the figure above that shows R=X@W and imagine single-head attention, this becomes clear. The number of elements in X, W, and R remain the same as with eight heads, but the way we interpret the elements in R changes. With a single head, we have just one query, key and value projection per token embedding (a row in P) and they would span one third of each row: 512 elements — the same as the embedding size.

What’s the point of multiple heads then, you might wonder? Well, Vaswani et al. argue that it allows heads to capture different “representation subspaces”. For example; one head might track syntactic relations, while another focuses more on positional information. There’s quite some work that investigates whether this indeed happens in practice, e.g. in translation. In fact, I did some work on this myself a few years ago in summarization.

3. Scaling in dot-product attention avoids extremely small gradients

Similar to the previous point, I never really thought about why we divide attention logits by some constant (🤖) but it’s actually pretty straightforward.

Recall that each logit is the result of a dot product (i.e. sum over the element-wise product) between a query and a key vector. A higher number of dimensions qkv_dim thus results in more products in that sum — causing higher variance in attention logits. As we can see in the examples below, a softmax transformation on logits with high variance results in extremely small output probabilities — and therefore tiny gradients.

4. The source embedding, target embedding AND pre-softmax linear share the same weight matrix

We now move away from multi head attention and dive into “weight tying” — common practice in sequence to sequence models. I find this quite interesting because embedding weight matrices actually make up for a huge number of parameters relative to the rest of the model. Given a vocabulary of 30k tokens and an embedding size of 512, this matrix contains 15.3 million parameters!

Imagine having three such matrices: one that maps source token indices to embeddings, one that maps target tokens to embeddings, and one that maps each of the decoder’s topmost contextualized token embeddings to logits over the target vocabulary (the pre-softmax linear layer). Yeah; that leaves us with 46 million parameters.

In the code you can see that I initialize one embedding layer in the main transformer class (🤖) that I use as encoder embedding (��), decoder embedding (🤖) and decoder pre-softmax transformation weights (🤖).

5. During inference we perform a decoder forward pass for every token; during training we perform a single forward pass for the entire sequence

This one might be obvious to some — especially those working on sequence to sequence tasks — but crucial to understanding how a transformer is actually trained.

Inference

Let’s say that we are in inference mode, in which we autoregressively (one by one) predict target tokens. The transformer always outputs a distribution over the vocabulary for each token in the batch. The next token is predicted based on the output distribution of the last token index in the batch (🤖). This means that we basically throw away all the output distributions for all the previous indices.

Training and Teacher Forcing

This contrasts training in which we use teacher forcing. During training, we perform just one forward pass through the decoder, regardless of the sequence length (🤖). We (the teacher) force-feed the entire batch of ground-truth target sequences, at once. This gives us all next-token predictions at once, for which we compute the average loss.

Note that each token prediction is based on previous ground-truth tokens and not previously predicted tokens! Note also that this single forward pass is equivalent to autoregressive decoding using only ground-truth tokens as input and ignoring previous predictions (!), but much more efficient. We use an attention mask to restrict the decoder self-attention module to attend to future tokens (the labels) and cheat.

I think it’s useful to realize that this way of training, called teacher forcing, is applied not only to translation models, but also to most popular pre-trained autoregressive language models like GPT-3.

6. Transformers can handle arbitrarily long sequences, in theory…

…in practice, however, multi-head attention has compute and memory requirements that limit the sequence length to around 512 tokens. Models like BERT do, in fact, impose a hard limit on the input sequence length because they use learned embeddings instead of the sinusoid encoding. These learned positional embeddings are similar to token embeddings, and similarly work only for a pre-defined set of positions up to some number (e.g. 512 for BERT).

7. Transformers are residual streams

On to the final point. I like to think of a transformer as multiple “residual streams”. This is similar to how an LSTM keeps a left-to-right, horizontal “memory stream” while processing new tokens one by one, and regulates information flow with gates.

In a transformer, this stream doesn’t run horizontally across tokens, but vertically across layers (e.g. encoders) and sub-layers (i.e. multi-head attention and fully connected layers). Each sub-layer simply adds information to the residual stream using residual connections. This is an awesome blogpost that discusses residual streams in more detail, and this is a cool paper that exploits the notion of “vertical recurrence”.

A consequence of this residual stream is that the number of dimensions of intermediate token representations must be the same throughout all (sub-)layers, because the residual connections add two vectors (example in the encoder: 🤖). On top of that, because the encoder (and similarly decoder) layers are stacked on top of each other, their output shape must match their input shape.

Wrapping up…

Thank you for reading this post! Let me know if you liked it, have questions, or spotted an error. You can message me on Twitter or connect on LinkedIn. Check out my other blogposts at jorisbaan.nl/posts.

Even though I wrote the code to be easily understandable (it’s well documented, unit tested and type checked) please do use the official PyTorch implementation in practice 😁.

Thanks to David Stap for the idea to implement a transformer from scratch, Dennis Ulmer and Elisa Bassignana for feedback on this post, Lucas de Haas for a bug-hunting session and Rosalie Gubbels for creating the visuals. Some great resources I looked at for inspiration are the PyTorch transformer tutorial and implementation; Phillip Lippe’s transformer tutorial; Alexander Rush’s Annotated Transformer and Jay Alammar’s Illustrated Transformer.

Implementing a Transformer From Scratch was originally published in TDS Archive on Medium, where people are continuing the conversation by highlighting and responding to this story.

Table of Contents

1. Preamble
2. Neural Network Generalization
3. Back to Basics: The Bayesian Approach
    3.1 Frequentists
    3.2 Bayesianists
    3.3 Bayesian Inference and Marginalization
4. How to Use a Posterior in Practice?
    4.1 Maximum A Posteriori Estimation
    4.2 Full Predictive Distribution
    4.3 Approximate Predictive Distribution
5. Bayesian Deep Learning
    5.1 Recent Approaches to Bayesian Deep Learning
6. Back to the Paper
    6.1 Deep Ensembles are BMA
    6.2 Combining Deep Ensembles With Bayesian Neural Networks
    6.3 Neural Network Priors
    6.4 Rethinking Generalization and Double Descent
7. Final Words

1. Preamble

Bayesian (deep) learning has always intrigued and intimidated me. Perhaps because it leans heavily on probabilistic theory, which can be daunting. I noticed that even though I knew basic probability theory, I had a hard time understanding and connecting that to modern Bayesian deep learning research. The aim of this blogpost is to bridge that gap and provide a comprehensive introduction.

Instead of starting with the basics, I will start with an incredible NeurIPS 2020 paper on Bayesian deep learning and generalization by Andrew Wilson and Pavel Izmailov (NYU) called Bayesian Deep Learning and a Probabilistic Perspective of Generalization. This paper serves as a tangible starting point in which we naturally encounter Bayesian concepts in the wild. I hope this makes the Bayesian perspective more concrete and speaks to its relevance.

I will start with the paper abstract and introduction to set the stage. As we encounter Bayesian concepts, I will zoom out to give a comprehensive overview with plenty of intuition, both from a probabilistic as well as ML/function approximation perspective. Finally, and throughout this entire post, I’ll circle back to and connect with the paper.

I hope you will walk away not only feeling at least slightly Bayesian, but also with an understanding of the paper’s numerous contributions, and generalization in general ;)

2. Neural Network Generalization (Abstract & Introduction)

If your Bayesian is a bit rusty the abstract might seem rather cryptic. The first two sentences are of particular importance to our general understanding of Bayesian DL. The middle part presents three technical contributions. The last two highlighted sentences provide a primer on new insights into mysterious neural network phenomena. I’ll cover everything, but first things first: the paper’s introduction.

An important question in the introduction is how and why neural networks generalize. The authors argue that

“From a probabilistic perspective, generalization depends largely on two properties, the support and the inductive biases of a model.”

Support is the range of dataset classes that a model can support. In other words; the range of functions a model can represent, where a function is trying to represent the data generative process. The inductive bias defines how good a model class is at fitting a specific dataset class (e.g. images, text, numerical features). The authors call this, quite nicely, the “distribution of support”. In other words, model class performance (~inductive bias) distributed over the range of all possible datasets (support).

Let’s look at the examples the authors provide. A linear function has truncated support as it cannot even represent a quadratic function. An MLP is highly flexible but distributes its support across datasets too evenly to be interesting for many image datasets. A convolutional neural network exhibits a good balance between support and inductive bias for image recognition. Figure 2a illustrates this nicely.

The vertical axis represents what I naively explained as “how good a model is at fitting a specific dataset”. It actually is Bayesian evidence, or marginal likelihood; our first Bayesian concept! We’ll dive into it in the next section. Let’s first finish our line of thought.

A good model not only needs a large support to be able to represent the true solution, but also the right inductive bias to actually arrive at that solution. The Bayesian posterior, think of it as our model for now, should contract to the right solution due to the right inductive bias. However, the prior hypothesis space should be broad enough such that the true model is functionally possible (broad support). The illustrations below demonstrate this for the three example models. From left to right we see a CNN in green, linear function in purple, and MLP in pink.

At this point in the introduction, similar to first sentence of the abstract, the authors stress that

”The key distinguishing property of a Bayesian approach is marginalization instead of optimization, where we represent solutions given by all settings of parameters weighted by their posterior probabilities, rather than bet everything on a single setting of parameters.”

The time is ripe to dig into marginalization vs optimization, and broaden our general understanding of the Bayesian approach. We’ll touch on terms like the posterior, prior and predictive distribution, the marginal likelihood and bayesian evidence, bayesian model averaging, bayesian inference and more.

3 Back to Basics: The Bayesian Approach

We can find claims about marginalization being at the core of Bayesian statistics everywhere. Even in Bishop’s ML bible Pattern Recognition and Machine Learning. The opposite to the Bayesian perspective is the frequentist perspective. This is what you encounter in most machine learning literature. It’s also easier to grasp. Let’s start there.

3.1 Frequentists

The frequentist approach to machine learning is to optimize a loss function to obtain an optimal setting of the model parameters. An example loss function is cross-entropy, used for classification tasks such as object detection or machine translation. The most commonly used optimization techniques are variations on (stochastic) gradient descent. In SGD the model parameters are iteratively updated in the direction of the steepest descent in loss space. This direction is determined by the gradient of the loss with respect to the parameters. The desired result is that for the same or similar inputs, this new parameter setting causes the output to closer represent the target value. In the case of neural networks, gradients are often computed using a computational trick called backpropagation.

From a probabilistic perspective, frequentists are trying to maximize the likelihood p(D|w, M}). In plain english: to pick our parameters w such that they maximize the probability of the observed dataset D given our choice of model M (Bishop, Chapter 1.2.3). M is often left out for simplicity. From a probabilistic perspective, a (statistical) model is simply a probability distribution over data D (Bishop, Chapter 3.4). For example; a language model outputs a distribution over a vocabulary, indicating how likely each word is to be the next word. It turns out this frequentist way of maximum likelihood estimation (MLE) to obtain, or “train”, predictive models can be viewed from a larger Bayesian context. In fact, MLE can be considered a special case of maximum a posteriori estimation (MAP, which I’ll discuss shortly) using a uniform prior.

3.2 Bayesianists

A crucial property of the Bayesian approach is to realistically quantify uncertainty. This is vital in real world applications that require us to trust model predictions. So, instead of a parameter point estimate, a Bayesian approach defines a full probability distribution over parameters. We call this the posterior distribution. The posterior represents our belief/hypothesis/uncertainty about the value of each parameter (setting). We use Bayes’ Theorem to compute the posterior. This theorem lies at the heart of Bayesian ML — hence the name — and can be derived using simple rules of probability.

We start with specifying a prior distribution p(w) over the parameters to capture our belief about what our model parameters should look like prior to observing any data.

Then, using our dataset, we can update (multiply) our prior belief with the likelihood p(D|w). This likelihood is the same quantity we saw in the frequentist approach. It tells us how well the observed data is explained by a specific parameter setting w. In other words; how good our model is at fitting or generating that dataset. The likelihood is a function of our parameters w.

To obtain a valid posterior probability distribution, however, the product between the likelihood and the prior must be evaluated for each parameter setting, and normalized. This means marginalizing (summing or integrating) over all parameter settings. The normalizing constant is called the Bayesian (model) evidence or marginal likelihood p(D).

These names are quite intuitive as p(D) provides evidence for how good our model (i.e. how likely the data) is as a whole. With “model as a whole” I mean taking into account all possible parameter settings. In other words: marginalizing over them. We sometimes explicitly include the model choice M
in the evidence as p(D|M). This enables us to compare different models with different parameter spaces. In fact, this comparison is exactly what happens in the paper when comparing the support and inductive bias between a CNN, MLP and linear model!

3.3 Bayesian Inference and Marginalization

We’ve now arrived at the core of the matter. Bayesian inference is the learning process of finding (inferring) the posterior distribution over w. This contrasts with trying to find the optimal w using optimization through differentiation, the learning process for frequentists.

As we now know, to compute the full posterior we must marginalize over the whole parameter space. In practice this is often impossible (intractable) as we can have infinitely many such settings. This is why a Bayesian approach is fundamentally about marginalization instead of optimization.

The intractable integral in the posterior leads to a different family of methods to learn parameter values with. Instead of gradient descent, Bayesianists often uses sampling methods such as Markov Chain Monte Carlo (MCMC), or variational inference; techniques that try to mimic the posterior using a simpler, tractable family of distributions. Similar techniques are often used for generative models such as VAEs. A relatively new method to approximate a complex distribution is normalizing flows.

4. How to Use a Posterior in Practice?

Now that we understand the Bayesian posterior distribution, how do we actually use it in practice? What if we want to predict, say, the next word, let’s call y given an unseen sentence x?

4.1 Maximum A Posteriori (MAP) Estimation

Well, we could simply take the posterior distribution over our parameters for our model M and pick the parameter setting w^ that has the highest probability assigned to it (the distribution’s mode). This method is called Maximum A Posteriori or MAP estimation. But… It would be quite a waste to go through all this effort of computing a proper probability distribution over our parameters only to settle for another point estimate, right? (Except when nearly all of the posterior’s mass is centered around one point in parameter space). Because MAP provides a point estimate, it is not considered a full Bayesian treatment.

4.2 Full Predictive Distribution

The full fledged Bayesian approach is to specify a predictive distribution
p(y|D,x).

This defines the probability for class label y given new input x and dataset D. To compute the predictive distribution we need to marginalize over our parameter settings again! We multiply the posterior probability of each setting w with the probability of label w given input x using parameter setting w. This is called Bayesian Model Averaging, or BMA, we take a weighted average over all possible models (parameter settings in this case). The predictive distribution is the second important place for marginalization in Bayesian ML, the first being the posterior computation itself. An intuitive way to visualize a predictive distribution is with a simple regression task, like in the Figure below. For a concrete example check out these slides (slide 9–21).

4.3 Approximate Predictive Distribution

As we know by now, the integral in the predictive distribution is often intractable and at the very least extremely computationally expensive. A third method to using a posterior is by sampling a few parameters settings and combining the resulting models (e.g. approximate BMA). This is actually called a Monte Carlo approximation of the predictive distribution!

This last method is vaguely reminiscent of something perhaps more familiar to a humble frequentist: deep ensembles. Deep ensembles are formed by combining neural networks that are architecturally identical, but trained with different parameter initializations. This beautifully ties in with where we left off in the paper! Remember the abstract?

“We show that deep ensembles provide an effective mechanism for approximate Bayesian marginalization, and propose a related approach that further improves the predictive distribution by marginalizing within basins of attraction”.

Reading the abstract for the second time, the contributions should make a lot more sense. Also, we are now finally steering into Bayesian Deep Learning territory!

5. Bayesian Deep Learning

A Bayesian Neural Network (BNN) is simply posterior inference applied to a neural network architecture. To be precise, a prior distribution is specified for each weight and bias. Because of their huge parameter space, however, inferring the posterior is even more difficult than usual.

So why do Bayesian DL at all?

The classic answer is to obtain a realistic expression of uncertainty or calibration. A classifier is considered calibrated if the probability (confidence) of a class prediction aligns with its misclassification rate. As said before, this is crucial in real world applications.

”Neural networks are often miscalibrated in the sense that their predictions are typically overconfident.”

However, the authors of our running paper, Wilson and Izmailov, argue that Bayesian model averaging increases accuracy as well. According to Section 3.1, the Bayesian perspective is in fact especially compelling for neural networks! Because of their large parameter space, neural networks can represent many different solutions, e.g. they are underspecified by the data. This means a Bayesian model average is extremely useful because it combines a diverse range of functional forms, or “perspectives”, into one.

”A neural network can represent many models that are consistent with our observations. By selecting only one, in a classical procedure, we lose uncertainty when the models disagree for a test point.”

5.1 Recent Approaches To (Approximate) Bayesian Deep Learning

A number of people have recently been trying to combine the advantages of a traditional neural network (e.g. computationally efficient training using SGD & back propagation) with the advantages of a Bayesian approach (e.g. calibration).

Monte Carlo Dropout

One popular and conceptually easy approach is Monte Carlo dropout. Recall that dropout is traditionally used as regularization; it provides stochasticity or variation in a neural network by randomly shutting down weights during training It turns out dropout can be reinterpreted as approximate Bayesian inference and applied during testing, which leads to multiple different parameter settings. Sounds a little similar to sampling parameters from a posterior to approximate the predictive distribution, mh?

Stochastic Weight Averaging — Gaussian (SWAG)

Another line of work follows from Stochastic Weight Averaging (SWA), an elegant approximation to ensembling that intelligently combines weights of the same network at different stages of training (check out this or this blogpost if you want to know more. SWA-Gaussian (SWAG) builds on it by approximating the shape (local geometry) of the posterior distribution using simple information provided by SGD. Recall that SGD “moves” through the parameter space looking for a (local) optimum in loss space. To approximate the local geometry of the posterior, they fit a Gaussian distribution to the first and second moment of the SGD iterate. Moments describe the shape of a function or distribution, where the zero moment is the the sum, the first moment is the mean, and the second moment is the variance. These fitted Gaussian distributions can then be used for BMA.

Frequentist Alternatives to Uncertainty Representation

I have obviously failed to mention at least 99% of the field here (e.g. KFAC Laplace and temperature scaling for improved calibration), and picked the examples above in part because they are related to our running paper. I’ll finish with one last example of a recent frequentist (or is it…) alternative to uncertainty approximation. This is a popular method showing that one can train a deep ensemble and use it to form a predictive distribution, resulting in a well calibrated model. They use a few bells and whistles that I won’t go into, such as adversarial training to smoothen the predictive distribution. Check out the paper here.

6. Back to the Paper

By now we are more than ready to circle back to the paper and go over its contributions! They should be easier to grasp :)

6.1 Deep Ensembles are BMA

Contrary to how recent literature (myself included) has framed it, Wilson and Izmailov argue that deep ensembles are not a frequentist alternative to obtain Bayesian advantages. In fact, they are a very good approximation of the posterior distribution. Because deep ensembles are formed by MAP or MLE retraining, they can form different basins of attraction. A basin of attraction is a “basin” or valley in the loss landscape that leads to some (locally) optimal solution. But there might be, and usually are, multiple optimal solutions, or valleys in the loss landscape. The use of multiple basins of attraction, found by different parts of an ensemble, results in more functional diversity than Bayesian approaches that focus on approximating posterior within single basin of attraction.

6.2 Combining Deep Ensembles with Bayesian Neural Networks (Section 4)

This idea of using multiple basins of attraction is important for the next contribution as well: an improved method for approximating predictive distributions. By combining the multiple basins of attraction property that deep ensembles have with the Bayesian treatment in SWAG, the authors propose a best-of-both-worlds solution: Multiple basins of attraction Stochastic Weight Averaging Gaussian or MultiSWAG:

”MultiSWAG combines multiple independently trained SWAG approximations, to create a mixture of Gaussians approximation to the posterior, with each Gaussian centred on a different basin. We note that MultiSWAG does not require any additional training time over standard deep ensembles.”

Have a look at the paper if you’re interested in the nitty gritty details ;)

6.3 Neural Network Priors (Section 5)

How can we ever specify a meaningful prior over millions of parameters, I hear you ask? It turns out this is a pretty valid question. In fact, the Bayesian approach is sometimes criticised because of it.

However, in Section 5 of the paper Wilson and Izmailov provide evidence that specifying a vague prior, such as a simple Gaussian might actually not be such a bad idea.

”Vague Gaussian priors over parameters, when combined with a neural network architecture, induce a distribution over functions with useful inductive biases.” …

… ”The distribution over functions controls the generalization properties of the model; the prior over parameters, in isolation, has no meaning.”

A vague prior combined with the functional form of a neural network results in a meaningful distribution in function space. The prior itself doesn’t matter, but its effect on the resulting predictive distribution does.

6.4 Rethinking Generalization & Double Descent (Section 6 and 7)

We have now arrived at the strange neural network phenomena I highlighted in the abstract. According to Section 6, the surprising fact that neural networks can fit random labels is actually not surprising at all. Not if you look at it from perspective of support and inductive bias. Broad support, the range of datasets for which p(D|M) > 0, is important for generalization. In fact, the ability to fit random labels is perfectly fine as long as we have the right inductive bias to steer the model towards a good solution. Wilson and Izmailov also show that this phenomena is not mysteriously specific to neural networks either, and that Gaussian Processes exhibit the same ability.

Double Descent

The second phenomena is double descent. Double descent is a recently discovered phenomena where bigger models and more data unexpectedly decreases performance.

Wilson and Izmailov find that models trained with SGD suffer from double descent, but that SWAG reduces it. More importantly, both MultiSWAG as well as deep ensembles completely mitigate the double descent phenomena! This is in line with their previously discussed claim that

”Deep ensembles provide a better approximation to the Bayesian predictive distribution than conventional single-basin Bayesian marginalization procedures.”

and highlights the importance of marginalization over multiple modes of the posterior.

Final Words

You made it! Thank you for reading all the way through. This post became quite lengthy but I hope you learned a lot about Bayesian DL. I sure did.

Note that I am not affiliated to Wilson, Izmailov or their group at NYU. This post reflects my own interpretation of their work, except for the quote blocks taken directly from the paper.

Please feel free to ask any question or point out mistakes that I’ve undoubtedly made. I would also love to know whether you liked this post. You can find my contact details on my website, message me on Twitter or connect on LinkedIn. Check out my personal blog at jorisbaan.nl/posts for proper math rendering!

A Comprehensive Introduction to Bayesian Deep Learning was originally published in TDS Archive on Medium, where people are continuing the conversation by highlighting and responding to this story.