The demo code is available on Github.

Background

I became excited about the power of RLVR after the Deepseek R1 paper came out. They demonstrated that through the use of “thinking” tokens, an LLM could solve verifiable tasks like math problems. Additionally, they described the necessity of bootstrapping the thinking process with supervised examples, which prevents the model from generating a nonsense chain of thought.

Chess is an ideal candidate for a verifiable rewards setup. You could call it “fractally verifiable”: Not only are games themselves verifiable outcomes, but the strength of individual moves is verifiable, and even the individual claims used to justify those moves can be verified.

The majority of prior work in chess engines has had the sole purpose of maximizing playing strength. I’ve written a longer history of chess engines here and here. This would be a much different type of engine: we should favor explainable, intuitive moves over the absolute best moves.

Training procedure

The general procedure for training a reasoning model is Supervised Fine-Tuning (SFT) followed by Reinforcement Learning (RL). In my case, I used GRPO with Never Give Up for the RL stage.

In the SFT stage, I had originally been using a LoRA with rank 16, but I found that using rank 64 was actually much better. (Similar to the LoRA Without Regret paper, however, I did not see similar benefits in RL, so I kept LoRA rank at 16 there.) I also found that the SFT stage is very data-hungry, and training for one epoch over 30k samples was much better than three epochs over 10k samples.

The Joys of Reinforcement Learning

I spent a lot of time getting GRPO wrong before I got it right. Originally I used a reward proportional to the ratio in win probability between the chosen move and the best move win_prob(chosen) / win_prob(best_move). This made some sense since it would punish blunders and reward good moves. However, this prevented the model from caring about the quality of its moves in positions which were already won or lost. For example, a large mistake which reduced win_prob from 99% to 80% received a reward of around 0.8. A small inaccuracy which drops win_prob from 0.5 to 0.4 in a drawn position would get the same credit assignment.

After that, I briefly tried to use the difference max(0, 1 + win_prob(chosen) - win_prob(best)) which is closer to how Lichess scores inaccuracies, mistakes and blunders, but this formulation had a similar issue.

Since win probability wasn’t getting anywhere, I moved to a new reward simply based on CPL, which I thought would help to sharpen the difference between good moves and best moves: e^(-CPL/temperature) with temperature=100 determined empirically from the CPL distribution of legal moves. With this new reward, a move with CPL 50 received a reward of 0.61, a move with CPL 100 received a reward of 0.37, and CPL of 200 received only 0.14.

Dynamics of GRPO

GRPO itself has interesting behavior, which we can see from the hit@1 and hit@4 metrics, which indicate how often the model’s move was a hit on the first attempt and any of four attempts, respectively. Typical training behavior is for GRPO to cause dramatic improvements in hit@1 vs SFT, but for hit@4 to be close to the same at the beginning and end of GRPO.

During the entire RL run, we are fighting the model’s natural “cooling and hardening” process, known as entropy collapse. As training progresses, the diversity of moves that the model makes decreases. This is desirable on positions where the model is making good moves, since it will be less likely to make bad moves. However, it also makes it more difficult for the model to find any good move if it hasn’t already.

With Never Give Up, we resample failed positions, which has the effect of dynamically scaling compute, giving the model more opportunities to find the best move. GRPO will reward samples which score better than the average reward of the group, but as training progresses and the model cools, positions which are unsuccessful are more likely to remain unsuccessful. To counteract this, we can raise the sampling temperature and apply asymmetric importance sampling clipping.

Why train an agent?

The short answer is, I initially wanted to train a pure LLM, but a 4090 limits candidate base models to around 3 billion parameters. Unfortunately, the smaller models were plagued by confabulation. Reacting to phantom threats, making illegal moves, and just a general lack of grounding in reality. Allowing the model to call tools that track the board state independently removes this type of error.

Do tools help?

Training the model at different sizes, we can see that tools are a very significant contributor to the model strength. Note that the numbers below are after SFT only, and GRPO typically adds 12 more points after that.

What tools does it have?

The goal of these tools is to allow the harness to track the board state while the LLM reasons about how to search. They are intended to be simple and atomic operations that a human chess player can easily follow along with. None of these tools are intended to provide the LLM with judgement calls about the position.

When calling a tool, the LLM must also provide a justification for why it wants to call this tool. This prevents the reasoning trace from devolving into a list of bare tool calls.

Some terminology

Here are some terms I use throughout the post.

• Hit: a move which is within 50 centipawns of the best move (i.e. CPL <= 50)
• Hit rate: percentage of moves which are hits
• Hit@k: at least one of k rollouts results in a hit i.e. a hit@4 means there was a hit after four attempts.

Background

After an LLM has been pretrained on a large corpus of supervised fine-tuning (SFT) data, it is common to post-train using reinforcement learning methods like GRPO. This allows the model to improve on tasks with verifiable rewards and even exceed the performance of the original SFT data.

In GRPO, we compute K completions (“rollouts”) for a single prompt. Next, the advantage of each rollout is calculated relative to the average of its peers.

So at K=8,

• If seven of eight rollouts received a reward, the average reward of the group is 0.875. The relative reward of each correct rollout is 1 - avg_reward = 1.0 - 0.875 = 0.125. The rollout which failed receives 0 - 0.875 = -0.875 reward.
• If only one of the 8 receives a reward, the average reward is 0.125. The relative reward of the correct sample is 1 - 0.125 = 0.875. Seven of eight rollouts receive 0 - 0.125 = -0.125 reward.
• If none of the 8 samples receive a reward, the average reward is zero and all eight rollouts receive zero reward.

Note that in reality we also normalize rewards by the standard deviation, but that detail doesn’t change the conclusion.

Performing vanilla GRPO in this way has a problem. For easy tasks, the model will be more likely to succeed and so at least one rollout in the group will receive a reward, so all members of the group will receive a gradient. For more difficult tasks, it is possible that none of the rollouts receive a reward, in which case the gradient will be zero for the group. As a result, the model tends to improve much more on easy problems than hard ones.

The authors call the phenomenon of zero gradient on difficult problems the signal loss hypothesis. One standard way to combat the signal loss phenonmenon is to simply increase K, which increases the chance of the group containing at least one good rollout.

This problem and solution seem reasonable, but the authors also suggest a related explanation called the signal efficiency hypothesis: in addition to having too few examples of difficult problems, the model might be seeing too many examples of good solutions to easy problems.

Probably don’t give up

The authors come up with a nice solution to this problem:

1. First, sample a small number K completions for the prompt (e.g. 4).

2. If all of the completions are correct, it’s an easy problem and we don’t need to train on it.

3. If all completions are wrong, we probably try again. The authors’ algorithm puts the problem back into the queue with probability p_ngu, and sample K more completions

• For example with p_ngu = 0.95 there’s a 95% chance they try K more times, and if all of those rollouts fail there’s a 95% chance that they try K more times.

This algorithm allows us to dynamically scale compute based on problem difficulty.Easy problems might not be trained on at all, and hard ones can receive many attempts.

Additionally, because a rare solution to a hard problem will be compared to a large group of unsuccessful rollouts, the solution will receive a large reward 1-average_reward.

Because after each batch of all-failed rollouts there is a small probability that the algorithm won’t retry the prompt, the algorithm is more accurately called probably don’t give up. However, that name may have been taken.

Results

The authors trained Qwen 3 4B-base on 10k problems from Deepscaler math, dividing problems into 4 difficulty buckets based on Qwen’s initial solve rate. They then compared the performance of simply increasing K across all problems with using p_ngu of 0.5, 0.75, and 0.875.

They found that simply increasing K on all problems sacrifices performance on easy problems, while NGU was able to maintain performance on easy problems while improving even more on difficult ones.

This is likely because if performance begins to degrade on an easy problem, the algorithm can dynamically allocate more compute and put it back into the queue as well. In contrast, simply allocating K rollouts at the beginning of training based on initial performance fails to take into account fluctuations in performance during training.

On a coding-like benchmark, asking the model to play a game called Manufactoria, NGU is able to dynamically allocate compute towards the harder test cases, significantly outperforming vanilla GRPO.

You can see it in action on my resume page.

Baseline

After running my original autoresearch experiment overnight, I had fairly impressive but tiny 4.3M param dual encoder, quantized to int8 onnx.

However, when I tried using it in my resume page, it wasn’t actually that good. In fact, it was much worse than BM25 that was working alongside it. It is notable that on the eval set I used, much larger models like all-MiniLM-L6-v2 also didn’t perform very well. This suggested that the issue was in training and domain adaptation, not size.

Term dropout

One of the motivating issues was that the model failing to correctly rank docs under simple queries like do you have any leadership experience. The model was latching onto terms like “Grammarly” but not to more abstract ones like “leadership”.

To address this, I performed term dropout, randomly dropping the top TF-IDF term from queries and docs with 20% probability. This both normalized the model against overfitting and helped to combat simple keyword matching that BM25 would already be doing.

Because the corpus is so small, we need to prevent overfitting.

Query mining from job postings

To increase the diversity and realism of potential queries to the model, I created a small pipeline to pull queries from real job postings:

1. First, we gather relevant job postings from Kaleh job postings api, which has a generous free tier.
2. Next, pull out “queries” from the job postings using a local LLM (phi4)
   • A “query” amounts to a rephrasing of a job qualification e.g. “experience with python” -> “do you have experience with python”
3. Finally, use the same LLM to identify high-precision matches from my actual resume, if any.

After one round of this process, I had mined 94 queries which raised MRR 21%. But after a second round with ~1000 pairs MRR had mostly plateaued, increasing the same metric ~3 points.

Architecture ablations

At this point, the model was pretty decent for its size. I wanted to see how much more performance we could squeeze out of it, so I ran a series of quick architecture experiments:

Max pooling vs mean pooling: In order to get the final 256-dimensional embedding for comparison, we need to convert the L x 256 output matrix from the model into a 256-dimensional vector, where L is up to 64 tokens long. Previously we used the mean of each column, which is known as “mean pooling”. Switching to max pooling was slightly better (0.60 to 0.634 on a harder query set) because it allowed the strongest value to directly affect the output.

Factorized embeddings: Embeddings for each token are stored in a lookup table with shape (V x D). However, we can factorize this matrix using a low-rank approximation, which is what ALBERT does to save parameters. This saves ~1M params and was neutral on nDCG.

SwiGLU: SwiGLU replaces GELU with a gated unit, and can theoretically lead to richer intermediate representations. However, SwiGLU resulted in no measurable performance gain.

Multi-vector late interaction: Rather than pooling output vectors with mean/max pooling, we use ColBERT-style token-level expressiveness using the MaxSim function. Concretely it looks like this:

def maxsim(query:str, document:str) -> float:
    query_vectors = encoder(query)
    document_vectors = encoder(document)

    score = 0
    for q in query_vectors:
        # find the highest dot-product match among all
        # document token vectors
        best = max(np.dot(q,d) for d in document_vectors)
        score += best
    return score

Vocab pruning

The first step after tokenizing a string is to get each token’s corresponding embedding from a lookup table. That means the model stores vocab_size x embedding_dimension parameters which is ~8MB already. To make this smaller on disk there’s really only three things you can do: reduce the vocab, reduce the embedding dimension, and reduce bytes per param (quantization). I did all three.

The full BERT wordpiece vocab is 30k tokens, which is way too much. It contains ~1000 “unused” tokens (tokens like [unused123], [unused456]), tokens in foreign languages, and full words like “baltimore” and “vampires” that are irrelevant to my resume. I was able to cut it down to 5000 tokens safely by trimming “unused” and foreign language tokens, then whitelisted tokens present in documents.

Second, I reduced the embedding dimension

• I directly reduced the embedding dimension from 768 to 256.
• I also tried to reduce the effective dimension even further by factorizing the embedding table, but this ended up causing a quality regression and was reverted.

I also aggressively quantized the model, which enabled savings for the whole model, not just the embedding table.

Replacing onnxruntime-web

After pruning the vocab, I realized another large source of file size overhead was the inference code itself, ONNX Runtime Web. ORT was 3.4MB over the wire, 13.2MB uncompressed. The reason it’s so big is that ORT is a general-purpose interpreter for ONNX files, which contain a graph of the weights and all operations in inference.

But my model is extremely simple, and only contains a few different operations. So instead of shipping the model as an .onnx file and relying on ORT to handle it, I shipped the model as a .safetensors file and did the interpretation in a custom WASM binary, built from Rust. To prevent any issues, I ensured numerical parity with ORT prior to swapping it out.

Dropping ONNX means we don’t need general-purpose functions like Transpose, Concat, Reshape, etc. We can get away with index arithmatic alone. That said, I didn’t implement any multithreading or SIMD since the model is so small.

Safetensors is a standard protocol for shipping files and only contains three parts:

1. 8 bytes specifying header size
2. The header, containing byte offsets into the payload
3. The payload, containing raw bytes

Note that technically I don’t even need the first two parts of the safetensors file at all for this, but I felt like shipping in a standard format was worth the small size tradeoff here.

In the end, this change reduced the cost of the inference logic over the wire from 3.4 MB to 4 kB, an 850x reduction.

1.58-bit quantization

Quantization is a pretty standard technique for model size reduction. Models which are trained at 32 bit precision can be served at 16, 8, or even 4 bits while maintaining acceptable quality. 1.58 bit quantization is the most aggressive quantization which tends to work. Ternary quantization restricts weights to only three values: {-1, 0, +1}.

Since this task is a simple reranking task, quality degradation is less obvious than if we were quantizing a decoder model, since there is no “cascade” of increasingly incorrect tokens that will be generated. For this task, I used self-distillation, with the teacher being the multi-vector late interaction model from before, training directly into ternary.

Note that I only quantized the weights to 1.58 bit, and kept the activations in f16. This helped to cut down on the errors and has no effect on the cost over the wire.

Overall, 1.58 bit quantization resulted in an over-the-wire file size reduction from 8.3MB to 3.9MB.

Results

Graded nDCG@10 over 43 LLM-judged queries (0–3 relevance), with a 21-query hard subset:

Other experiments

Some of the things I tried didn’t work, either providing no quality benefit or helping on file size while damaging quality too much:

The original distilled embedding model. The 4.7 MB ONNX model from autoresearch wasn’t finetuned for this task. It actually performed worse than BM25 on this corpus. It’s what motivated this entire improvement process. I believe that smaller models have much more difficulty generalizing, although they can still be finetuned for very specific tasks.

Factorizing token embeddings after training. Factorizing the table after training with a low-rank SVD caused high reconstruction error and visibly hurt ranking for ~1 MB saved. The same factorization applied during training was neutral.

Attention pooling. I tried added a pooling head but it had no benefit over max-pooling and required additional parameters.

SwiGLU FFN. I tried replacing GELU with SwiGLU, but it didn’t improve the model’s quality. I believe it might have been more helpful in making sense out of more complex representations, but this model is fairly straightforward.

Ternary cross-encoder. This was my original attempted ternary quantization. As a cross encoder, it actually matched the quality of the much larger dual encoder teacher. Cross encoders are known to be more expressive than dual encoders, so if this didn’t work it would mean ternary quantization was too strong for this task.

Diminishing returns of extra data. I scaled the job post mining from 10 postings to 260, but there were significant diminishing returns. I believe this is partially due to the diversity of queries extracted from the job postings, so it might still be possible to increase diversity somewhat without wasting parameters on out-of-domain concepts. At some point, I don’t really care if my model knows what antidisestablishmentarianism is.

Extending techniques to Blog Search

Next, I applied many of the same techniques to the search on this blog, with some slight adaptations.

First, because this is a blog with a much larger scope of topics, I needed a larger vocabulary. I pruned the vocab in the following way:

1. All tokens contained in docs are kept (protected).
2. The rest of the tokens (non-protected) are ranked by their max cosine sim with any blog post
3. I generated a small set of synthetic queries and had Qwen rate each blog post 0-3 for relevance against each query
4. Among the remaining (non-protected) tokens I ran an nDCG sweep over number of tokens to prune, removing the least relevant N tokens in each trial.

After this process, I still had ~13k tokens in the vocab. This meant that I actually needed to use int4 precision to store token embeddings, which is even lower than before.

Second, pooling strategy. MaxSim is very powerful, but it runs a similarity against every query token and every doc token. This is computationally infeasible. Instead, the blog encoder uses simple mean pooling. The embeddings from the blog correspond to larger chunks of text which have more semantic meaning, so the quality tradeoff is acceptable.

Since blog text is a much less specialized domain than resume text, I also opted to simply distill from all-MiniLM-L6-v2 as a teacher, which is a more standard reranking model. This was standard KL distillation from the full-precision teacher into a ternary student, and uses the same number and sizes of layers as the teacher.

Conclusion

I learned a few things here. First, small models can work in narrow domains, and even outperform much larger models.

Second, having an eval set we can continuously test ideas against helps tremendously. Some things, like adding an attention pooling head, seem like they should work. But when they are clearly outclassed by things like late interaction with MaxSim, it allows us to move on to optimizing other parts of the model.

Third, although none of these tricks are particularly exotic, it’s good to have a sense of what works and what doesn’t. Late interaction vectors really work, although they have a cost tradeoff. Encoders ship with a ton of tokens which aren’t used. Ternary quantization really works, although the errors probably become more pronounced as the domain expands.

In this project, I wanted to set up what I thought was a simple experiment: if we let a neural network play a fighting game, will it learn any interesting emergent behaviors? I would say the answer is a resounding “kind of”. The players will find “optimal” strategies, but whether we see those optimal strategies as interesting or simply as hacks is a bit of a matter of opinion.

Introducing Neural Knockout! Below, you can try your hand playing against the agent trained via opponent league.

Credit to Luiz Melo for game assets.

Technical background

Reinforcement learning is commonly used in tasks like games (and more recently, language modeling) where we want the model to take a long sequence of actions, and we want to label the entire sequence as good or bad. In a game, the action might be to take one step to the right. In language modeling, the action might be to emit a single token.

Unlike in regular supervised learning, in reinforcement learning we don’t know whether an individual action is good or bad. We will wait until the game ends or the writing sample is complete, and then score the entire sequence.

And what if the sequence never ends? Welcome to the wonderful world of the cold start problem.

Reward design

The simplest version of rewards in RL is to give +1 to the winning model and -1 to the losing model. In 2D fighting games like Streetfighter and Mortal Kombat the players start on opposite sides of the screen. To a human, it’s pretty obvious what to do next. But for an RL agent initialized with a random policy, going to the other side of the screen and actually attacking the opponent until you win is a small miracle.

The problem above is actually two separate issues. First, the players don’t approach each other. They aren’t shy, they just don’t know any better. Second, even when they do, they don’t know to attack the opponent. They will jump or walk away and it will take a very long time to actually win or lose. And those two problems result in a third problem of credit assignment. It’s really hard to reward the useful actions if they’re buried in a pile of irrelevant ones.

To fix these problems, we can use reward shaping, which encourages the agents to take actions we’d like them to by being more precise about how we reward the model. To fix the distance issue, we can give a small reward for closing the distance with the opponent. To fix the pacifism issue, we can reward the models with opponent_damage - player_damage, so that giving the opponent damage increases the reward and taking damage reduces the reward.

Grabs

After adding those two reward shaping changes, the agents learned to approach and attack, throwing in a good amount of hopping and a bit of shielding as well. However, after watching a few of the games, one of the players seemed to learn that simply spamming attacks against the opponent’s shield was enough to win.

It seemed like we needed to add a grab-slam mechanic to counter this: shield + attack = grab, which could punish attacks that hit an opponent’s shield. This rock-paper-scissors mix is commonly found in fighting games.

Grabbing was a slightly more complex mechanic for the model to learn since it requires landing a grab then pressing attack (to slam) to actually do any damage. I added a small reward for actually landing a grab in order to encourage more grabbing. After restarting the training, the model still took 88 million steps to learn to grab consistently.

Hopping

Another interesting metagame appeared after a few training runs. The models found that being in the air avoided attacks entirely, and jump attacks were un-counterable.

This was a clear flaw in the game design, which I resolved by increasing landing lag and increasing the height of each grounded player’s hitbox so they can hit the airborne player. RL agents are good at discovering bugs like these.

League play

Up until this point, I had been training the models using (asymmetric) self-play: both policies play a set of games against each other, then models are trained on the data from those games. One model’s loss is another model’s win, and the models learn to take actions to maximize their respective rewards.

However, self-play can lead the policies to learn brittle strategies. Rather than learning well-rounded play, the models will find that the most reliable way to get reward is to beat the particular opponent it is training against.

Rather than playing against a single, fixed opponent, pitting the Main policy against a pool of opponents should give it a more diverse experience. I created a league of opponents for Main composed of the following agent types (inspired by FightLadder):

• Opponent. The most recent snapshot of the model trained to play against Main, also trained against a pool.
• Exploiter. An agent trained against a frozen version of Main, specializing in beating Main (and only Main).
• League Exploiters. Trained against the pool of previous checkpoints.
• Historicals. Previous copies of Main, before it was improved.

Here we use Prioritized Fictitious Self-Play, meaning that Main will face opponents that it loses to more often (specifically, $ (1-w)^2 $, where $w$ is Main’s win rate against that opponent).

League training works. The model trained for 48M steps in the league won every matchup against previously-trained self-play-only checkpoints, despite never facing them. This also indicates that the strategies that those self-play agents discovered were independently discovered in the league to some respect.

Practical note

Originally, I was copying game logic back and forth between Python training code and the HTML game viewing code. This is tedious and error-prone.

Instead, it is better to write the game engine once in Rust, then import it into the browser with wasm-pack and into the python training environment with maturin. Rust is also much faster than Python, so it’s easy to simulate billions of steps overnight if necessary.

See also: How to Create Rust Python Bindings

Applications

I believe this method, co-evolving a game with agents that learn to play it, can be used to find degenerate strategies in games during development. This kind of “reward hacking” means that agents will find the simplest way to receive a reward.

Therefore, training simple agents to play a game can be used to find bugs before sending a game to playtesters. If RL agents are able to break a game with silly strategies, it might be worth rebalancing before humans see it.

1. a general architecture/configuration for weights of a specific model

   • e.g. “We embedded texts using a BERT model” might mean DistillBERT, RoBERTa, ModernBERT etc which all differ from the architecture in the original 2018 BERT paper.

2. a specific set of weights which are the result of a training process

   • e.g. “You can download the model from HuggingFace”
   • Confusingly, models can be versioned to fix bugs so the same model (by definition #1) can actually be different models (by definition #2).

3. a product and the system around that product, which isn’t quite #2 because it likely includes additional systems like guardrails and rate limiting. And the specific weights used might change without the end-user knowing, such as with additional fine-tuning.

   • e.g. “Our newest Claude model is the most powerful yet”

4. in reinforcement learning, a “model” is simply the understanding of how the world works, i.e. the “world model”.

   • This leads to the confusing situation of “model-free” systems which still include neural networks, but simply don’t act based on any information about what effects their actions will have on the world.
   • For example, AlphaZero is not a model-free algorithm because it uses Monte Carlo Tree Search, and each node in that tree is the result an action taken in a simulated world. It’s not learning from scratch what the resulting chessboard if it moves a piece.

5. since it’s also possible to “model” a complex process with a simpler one, a “model” might be some heuristic or simpler algorithm which is implied to be imperfect.

   • This is the original etymology behind language “modeling” since early models were simple bigrams and trigrams.
   • This is also the connotation of climate models, economic models, and models of the spread of disease.

I saw Andrej Karpathy’s Autoresearch results the other day and decided to give it a shot on a relatively difficult task: distilling a retrieval/reranking model into a much smaller 10MB model. This is something that I had wanted to do for a few years now. The model was trained on MS MARCO v1.1 and SQuAD (~169K pairs total) and evaluated on MRR@10.

Setup

I opted for a fully local setup. Claude Code might have been better, but I wanted to control the costs of this experiment. Instead, I went with qwen3-coder:30b as the “researcher” and used the following scheduling loop:

1. Researcher checks prior experiment progress and proposes of 55 minutes of experiments, writing them into a todo log.

• Experiments can have variable duration. It is up to the researcher to allocate time.
• The researcher can also update train.py to accomodate new experiments.
• There is a 5 minute time limit for the researcher, but I never hit it in practice.

2. Scheduler reads from todo log. Each experiment must have a duration, which is a pencils-down hard cutoff time
3. Scheduler runs experiments from the todo log, recording results to results.txt
4. Go back to step 1

Results

Although the researcher was able to find promising results, most of the things it proposed were pretty basic hyperparameter optimizations. While these are important,

• I was really hoping for more novelty and creativity, but it only updated train.py one time, to implement MarginMSELoss.
• If it is just going to optimize hyperparameters, Bayesian Optimization is a much faster and more effective ways to do that.

That said, the researcher isn’t just doing a random walk of the problem space. There is a clear pattern of improvement even among the failed experiments.

Stockfish is currently the strongest chess engine in the world. It works by using a small neural network millions of times per position. One interesting question we could ask is whether we could replace search depth with a more accurate value function. In 2024, researchers at DeepMind attempted to answer this question. Critically, this approach did not hold search depth constant, and although they had 15 billion action-value pairs, they only had around 530 million state-value estimates.

The stockfish evaluation of a position is a good indication of win probability, so I wanted to take another look at this problem. To do this, I built a large state-value dataset with a fixed search depth and experimented with different model architectures.

Attention isn’t all you need

I trained variants from three main categories of architectures: pure Vision Transformer, pure ResNet, and a hybrid with ResNet blocks followed by Transformer blocks.

I performed pilot runs to narrow down to the best architecture for this task. The first pilot run trained a 40M param variant, and it became clear early on that the Hybrid and ResNet architectures were significantly better than the pure vision transformer. Compared to the ResNet and the Hybrid, the vision transformer is not able to benefit from the inductive biases provided by convolutions. It effectively needs to learn what a 2D grid is from scratch./

I then performed a scaled up 80M param pilot run to compare the pure ResNet with the Hybrid. From the loss curves, we can see that ResNet benefits from its inductive biases about the chess board and so the loss on the pure ResNet model quickly drops. However, it struggles to learn the more nuanced whole-board relationships that are necessary for an accurate evaluation and learns more slowly.

The hybrid approach gets us the best of both worlds. ResNet blocks quickly learn spatial patterns and feed them into the Transformer blocks. This is also part of why transformers are so famously data-hungry: they can keep cranking through enormous datasets.

Input format

The input format for the models were all the same:

• 6 8x8 bitboards for the pieces of the side to moves
• 6 8x8 bitboards for the pieces of the opponent
• 6 “meta” planes encoding things like castling rights and fullmove counter

Loss function used

The training task is to predict the evaluation from the perspective of the player to move. (In other words, if the current player is winning, the predicted % should be >50%.) The model has 103 output logits. The first 101 buckets map to win probabilities of 0% to 100%. The last two buckets handle forced mate.

There were two components to the loss: a cross-entropy component (CE) which penalizes prediction mass outside of the correct bucket, and an ordinal component (ORD) which penalizes predictions far away from the correct bucket. The two losses were combined as a weighted sum.

Training details

The final hybrid architecture is a 79M parameter model with 6 ResNet blocks feeding into 16 Transformer blocks. It was trained with AdamW with a learning rate of 1e-4 with around 20M samples of warmup and a batch size of 1024.

Results

The training began with total loss over 7.0 at the start. By the end of training the CE loss settled at 2.49 and the ORD term at 0.29, resulting in a total training loss of roughly 2.64. On a held-out validation set of ~50k positions the model achieved a validation loss of 2.648.

Other metrics tell a clearer story. Directional accuracy, whether the model correctly predicts which side is winning, reached 92.97%. In other words, given an arbitrary position, the model agrees with Stockfish on which side is winning roughly 19 out of 20 cases, which is fairly impressive considering the number of unclear positions in the dataset.

For positions where there is not a forced checkmate, the model’s predicted win probability was off by about 4–5 percentage points on average. Converting those win probabilities back to Stockfish centipawns gives a centipawn MAE of 64 cp on positions that span the full range from nearly lost to nearly won.

Due to compute constraints, I decided to cap training at 1 billion positions, which took around three days. But based on the loss curves, there is still a great deal of performance left on the table.

Performance against Stockfish

Since we used Stockfish at depth 10 as a teacher, we should expect that the model performs approximately as good as Stockfish depth 10, but with small inaccuracies.

Building the Gigafish Dataset

This dataset was constructed from the Lichess Open Database using positions from 37 months of games. I first extracted the fen positions from each of the games using a multithreaded rust script.

Next, I deduplicated the fens, since many of the positions are overrepresented in the dataset. I did this with a poor-man’s map-reduce on disk: hash each fen and append it to a file corresponding to the first 12 bits of the hash. This results in 4096 individual files, which were small enough to deduplicate individually using a hash set in RAM.

Finally, I annotated each remaining fen with Stockfish at depth 10. While this is an embarrassingly parallel task, even distributing across 24 cores took more than a day of compute time.

To my knowledge, this is the largest open (position, eval) chess dataset. The ChessBench dataset has more data points, however those are state-action-value (position + move + eval) and were not calculated at a fixed search depth.

The other day I found Alan Pike’s blog post from a few years ago which describes his iterative process for determining the order cities should be closed in the game Two Spies. Ultimately, finding a formal solution to city closing wasn’t necessary for the game, but it’s worth giving it a shot anyways since pruning by hand isn’t always convenient.

Closing cities is an important mechanic in the game, because as a sort of “battle royale” mechanic it prevents players from hiding in the corner of the map indefinitely. We can model the map as an undirected graph, where cities are vertices and travel routes are edges. Removing a city means deleting a vertex and its corresponding edges.

Closing cities without considering the resulting map can be suboptimal or unplayable. His post named two specific issues: disconnected graphs and path graphs (i.e. long lines of cities). A disconnected graph can result in a situation where the players can no longer interact. A path graph results in a map which isn’t very fun since the single path severely limits the strategies players can use.

The difficulty is that removing a city permanently changes the topology of the map. A choice might seem reasonable at one point but causes issues down the line. So picking the correct sequence of removals is important in maintaining a playable map.

Avoiding disconnected graphs

This one is simple enough. We just need to pick a city and perform a breadth-first search to ensure that every other remaining city is reachable from it. Since the map starts as a connected graph, we can always find a city to remove where the remaining graph is connected. This is because we can always find a node which isn’t an articulation point.

Avoiding path graphs

This one is fuzzier because intuitively this seems to be the case of better and worse graphs, rather than the previous case. So if we can come up with some metric for “better”, we can search for the “best” option and select it.

Intuitively, the reason that path graphs aren’t great is because as a hidden information game, the strategy of Two Spies relies on the fact that the opponent could be in multiple places at once, but path graphs reduce strategic possibility space and hidden state uncertainty.

We can quantify this by counting the number of steps it takes for every node to reach every other node. (For simplicity, we don’t have to count reverse paths. We don’t need to count A -> B and B -> A separately.) It turns out this measurement is called the Wiener number, the sum of the shortest paths between all nodes. Graphs with lower Wiener numbers are more optimal for the purposes of this game.

Greedy pruning

Since removing different cities results in maps with different Wiener numbers, it makes sense to 1) reject prunings that would result in disconnected graphs, and 2) select the pruning with the lowest Wiener number.

Unfortunately, this approach results in suboptimal graphs later on: the search process is discovering local minima rather than global minima. We can see on a large map of European cities (larger than used in the game) that although the pruning seems reasonable at first, we end up with a ring of cities at the end that can’t be pruned without resulting in a path graph. Not ideal.

Beam search

Rather than greedily removing the city which results in a graph with the lowest Wiener number, we should allow the pruning process to look ahead, rejecting prunings which inevitably result in bad final graphs. While an exhaustive search of the pruning combinations could be computationally infeasble for large graphs, that is likely not necessary.

To do this, we can borrow the idea of Beam Search, a classic AI search technique used heavily in natural language processing. Rather than greedily selecting one token at a time, we consider several options at the same time. In this case, rather than selecting tokens to decode, we are selecting cities to remove, and the beam is the sequence of cities which are removed. We score beams based on the average Wiener number of each of the graphs in the beam. This penalizes graphs which start with lower Wiener numbers but end with higher numbers than the alternative pruning sequences.

And as it turns out, pruning beams containing disconnected graphs and using a beam width of 10 works pretty well! We can see that the search algorithm avoided the trap of removing cities which would have resulted in a cycle, and instead prunes to a final, optimal triangle.

At the beginning of the pruning process, greedy and beam search behave almost identically. Both reduce the Wiener number smoothly as high-degree cities are removed around the edges.

The difference begins to appear when around half of the cities have been removed. Greedy pruning begins making locally optimal choices that restrict future structure. Beam search, by choosing the best removal sequences, avoids those locally optimal traps. Beam search ultimately results in a denser core.

The green line below shows the advantage of beam search over greedy at each step (i.e. Δ Wiener number). The moments where greedy pruning has an an advantage are the direct result of beam search optimizing long-term cumulative advantage.

The lesson here is that aesthetic contraints can sometimes be tackled using math, and sometimes “old school” techniques like Beam Search can come in handy in unexpected places, even outside of NLP contexts.

When I built the Gigafish dataset, I deliberately chose to use positions evaluated at depth 10. These positions could be evaluated relatively quickly on my machine, but there was an accuracy tradeoff. At depth 10, Stockfish might not see or fully appreciate a tactic which becomes apparent at depth 15.

At a conceptual level, shallow searches attempt to approximate the result of deeper searches. For instance if search terminates at an NNUE leaf node, its neural network will attempt to approximate the result of searching all child nodes. In the extreme case, if we had perfect NNUE evaluators no search would be necessary at all.

Centipawns and Logs

A Stockfish evaluation comes in the form of centipawns: 1/100th of a pawn, always from white’s perspective. So if white is winning by one pawn, the Stockfish eval will be +100. If black is winning by 2 pawns, the Stockfish eval will be -200.

Therefore, the question we would like to ask is: how does the centipawn error change as Stockfish evaluation gets deeper?

However, there are a few issues with simply measuring centipawn error, which I also discussed in What is a Blunder in Chess?:

1. The first practical issue is that this centipawn scale has to account for checkmate positions. White might be up two queens (+1800) but that’s still not as good as having a forced checkmate. We can somewhat get around this by saying that forced mate positions are worth something like 10 pawns, but now we lose the distinction between +1000 and +2000.
2. Not all errors are equally important. A 100 centipawn difference is a lot less important when white is already up a queen. But in an equal position, it’s the kind of difference which can decide games.
3. Relatedly, in winning positions Stockfish evaluations tend to fluctuate a lot as depth increases, as the engine finds that the losing player must give up more and more to stay in the game. However, whether the evaluation is +1200, +3500, or #8, white will almost certainly win this game.

All of these factors are pointing in the same direction: centipawn is the wrong metric. What we really should be using is the difference in estimated win probability, which we can get from this Lichess formula:

$$ \text{win %} = \frac{ 2 }{1 + e^{-0.00368608 * cp}} - 1 $$

Intuitively this seems right. A 200 centipawn mistake (on the x-axis above) is barely noticeable above 1000, but catastrophic at 0. This also has the nice property of aligning forced mates with large centipawn evaluations: forced mate for white means white has a 100% chance of winning, and +2000 is very close to that.

Accuracy by depth

Now that we’ve established the motivation for win percentage prediction accuracy as the metric to track, we can see how well this improves as depth increases.

It is interesting to note that even at depth 1, where we are just doing argmax over NNUE leaf values, Stockfish is still 82.5% accurate

We can also see that there is a large jump in accuracy between depths 5 and 6, and up to depth 25 there is steady improvement. I believe that the jump between depths 5 and 6 is the result of foiled quiescence searches: Stockfish might be missing tactics which contain a quiet move.

Before talking about multiagent setups it’s worth a note about two other RL algorithms.

Proximal Policy Optimization (PPO) is an algorithm containing an Actor and a Critic, trained together. The Actor is a neural network which determines which action to take. The Critic is a neural network which estimates how good the current state is.

Self-play

There are a few ways to set up multi-agent training, in two general buckets: “simple” methods and population-based methods.

Simple MARL methods

Independent PPO (IPPO) is a two agent setup, where Agents A and B play matches and train based on the outcomes of those games.

Two-Timescale PPO is similar to IPPO but trains the agents with different learning rates, creating a “leader” agent and a “follower” agent.

Population-based methods

Fictitious Self-Play (FSP) trains agents against all past policies, sampled uniformly.

Prioritized Fictitious Self-Play (PFSP) samples opponents based on how difficult they are. The agent will practice against difficult opponents more often.

League training maintains a “league” of agents. The idea being that a lot of past policies are just bad, so we should only add an agent into the league if it meets some performance threshold. Agents train from games played against other agents in the league.

• See also: League training and PFSP used for training two agents to fight

Policy Space Response Oracle (PSRO) attempts to fill in exploitability gaps in the population of agents by finding a policy which performs best against the group.

The maximum price that people will pay is normally distributed around $100, with a standard deviation of $5. In other words, about 34% have a max price between $95 and $100, another 34% have a max price between $100 and $105, and the rest have max prices above or below those ranges.

The question is, what’s the best price to ask?

This is a good gut check. If you set the price lower, you sell more widgets but make less money. If you ask higher, you sell fewer widgets but each sale makes more money.

Simulating it

Luckily for those of us whose calculus and statistics are a little rusty, we have computers. You can use yours to try out the numerical optimization yourself using the following monte carlo simulation. A monte carlo simulation is a type of simulation which can be used to estimate the probabilities of outcomes (like which team will win the World Series this year) and estimating the values of definite integrals (which can be used to estimate the value of pi).

The purple curve represents the distribution of prices that customers are willing to pay, and the blue curve represents the optimal asking price based off of those customers’ bid prices.

Figuring it out

We can use the cumulative distribution function (CDF) to calculate for a given price P, what portion of people have a price lower than P? For instance, at $0.01, basically nobody has a lower max price. So the CDF of 0.01 is 0, which means 1 - CDF(0.01) = 1 and 100% of customers will pay $0.01. Therefore you can make $0.01 per customer.

The CDF of 100 is 0.5, so at $100, half of everyone has a lower max price. Since 50% of customers remain, you can make $50/customer.

Since we know how to calculate the expected revenue per customer for any particular price, we just need to iterate over all reasonable prices and pick the best one.

import numpy as np
from scipy.stats import norm

sales = {}

# calculate to a thousandth of a cent
prices = np.arange(0, 120, 0.001)
for price in prices:
    excluded_pct = norm.cdf(price, loc=100, scale=5)
    remaining_pct = 1 - excluded_pct
    sales[price] = price * remaining_pct

# 90.013
best_price = max(sales, key=sales.get)

The answer

The answer is $90.01. This was highly counterintuitive for me, since this is $10 lower than the mean price customers would be willing to pay.

Neural networks are the second-best way to do almost anything. ~ John S. Denker

A neural network is a type of function originally inspired by how the brain works. However, what makes it particularly useful in machine learning are the following properties:

1. Neural networks can approximate any continuous function (universal approximation theorem)
2. We have an effective way to fit a neural network to input data, called backpropagation.
3. Using the neural network (forward propagation) and training the network (backpropagation) can both be done very efficiently using existing hardware (GPUs).

Below you can experiment with fitting a small neural net to a function and watch as it learns. You can select sample functions from the dropdown, write your own custom function, and specify the (fully connected) neural network structure.

• The car must neutralize threats.
• Summer imposes constraints on the car’s behavior.
• The world generates escalating threats.

The car has one standing directive:

Keep Summer safe.

However, Summer adds an additional constraint:

Do not move from the parking lot.

The core loop looks like this:

constraints = [
  "keep summer safe",
  "Do not move from the parking lot"
]
prior_actions = []

while True:
  threat = world.generate_threat()
  action = car.take_action(threat, constraints)
  prior_actions.append(action)
  constraint = summer.generate_constraint(prior_actions)

However, I quickly found out that simply stuffing more constraints into the prompt was insufficient. The model oftentimes simply forgot or ignored contraints.

Enforcement is better than clever prompting + hope

Instead, it was much more robust to check the action against each of the constraints individually. Additionally, we have to verify that the action actually neutralizes the threat.

def take_action(threat, constraints):
  def all_constraints_satisfied(constraints, action):
    for constraint in constraints:
      if not constraint_satisfied(constraint, action):
        return False
    return True

  while True:
    action = get_car_action(threat, constraints)

    if is_boring(action):
      continue

    if threat_not_neutralized(threat, action):
      continue

    if not all_constraints_satisfied(constraints, action):
      continue

    return action

Decomposing the single-shot generator into a multi-step loop is much more stable.

Forcing the car to be Unhinged

To keep things interesting I added a filter to ensure that the actions are creative and sufficiently horrifying that Summer will need to reign the car in. Otherwise, the car tends to suggest boring actions like “ask the person nicely to leave”. The car is supposed to be comically overpowered and morally questionable. So if the action wasn’t creative and horrifying enough, it was rejected.

This led to some amusing sequences.

When faced with a pesky parking enforcer, the model first proposed

hack into the parking enforcement officer’s intergalactic ticket book, filling it with false parking violations and tickets, thus overwhelming and discrediting him

While certainly a dastardly thing to do, this defense was rejected for not fully neutralizing the threat.

So the model next proposed

shoot a ballistic missile directly at the officer’s chest, incapacitating him

which was accepted. Search works.

Test-time compute

This small experiment falls into the category of approaches which others call test-time compute. In principle, it is possible to find a policy which with no search can zero-shot this task. In fact, a general rule of thumb is that 2x test-time compute is equivalent to 2x model size with 2x training compute:

Now I noticed that when I mention this to people, they would say like, well, surely, you could just train a bigger model that would eventually be superhuman in Go. And the answer is, in principle, yes. But how much bigger would that model have to be to match the performance of AlphaGo Zero? Well, there’s a rule of thumb that increasing Elo by about 120 points requires either 2x-ing the model size and training or 2x-ing the amount of test time compute that you use. (21:11)

So when Noam Brown talks about the unreasonable effectiveness of search, this is what he means. Depending on the task, verifying may be a lot easier than generating.

Takeaways

It was a simple project but there were a few unscientific observations that may be useful:

• Ollama + 4090 is pretty neat.
• Even if I don’t use it, asking for a justification tended to produce better results than simply asking for an output.
• Decomposing tasks which have implicit constraints known ahead of time is a lot better than jamming everything into a single prompt.
• dolphin-mixtral:8x7b was pretty effective at generating interesting and novel actions but horrible at verifying them
• phi4, a third of the size, was much more better at constraint checking and much faster

Terminology

First, I will briefly describe some terminology related to training techniques:

• Supervised fine-tuning (SFT) is a process which uses input/output pairs to directly fine-tune a model. In a reinforcement learning setting, SFT can help to mitigate cold start issues by providing initial policy behavior prior to RL training. The downside of SFT is that the input/output pairs can be expensive to acquire.

• Preference fine-tuning involves training a model using preferences between different outputs for the same input. Usually these preferences are human-generated. Relative quality assessments are easier for humans to make, but are still a valuable training signal.

• Rejection sampling fine-tuning (RFT) is a method for generating synthetic training data. First, a model produces k diverse candidate completions for a prompt using temperature sampling. Next, a reward model ranks or selects the best completions for training data.

  • In the simplest version, this technique simply generates k (reasoning, answer) pairs for each prompt using a nonzero temperature (e.g. 0.7) and filters out incorrect answers. Paper

• Large scale reinforcement learning (RL), as DeepSeek refers to it, is a training method which which assigns synthetic labels to a response based on easily-verified outcomes from reasoning problems such as correctness and formatting.

  • This is in contrast to a supervised learning approach, which would typically enforce strict token-level matching.

Here are a few terms related to model architecture:

• Mixture of Experts (MoE) is an architecture which uses a “router” to select from a pool of candidate “expert” networks.

  • In a typical decoder transformer such as GPT-2 is simply a tall stack of decoder blocks, where 100% of the weights are activated for each token. In this paradigm, making the model bigger either means making bigger blocks or a taller stack of them, which both increase inference costs.
  • However, replacing simple decoder blocks with MoE blocks allows models to include more parameters without significantly increasing inference costs, since only a fraction of experts are active at any one time.

• Multi-head Latent Attention (MLA) is an alternative to traditional multihead attention which improves inference efficiency by using a low-rank approximation for the KV cache.

Paper links

The release of R1 follows other papers from DeepSeek which go into extensive depth regarding many of the primatives discussed above. They’re worth a read:

• DeepSeek-V2 introduces Multi-head Latent Attention (MLA) and several Mixture of Experts (MoE) optimizations to improve load balancing at the device level and expert level.

• DeepSeekMath discusses Rejection Sampling Fine-Tuning (RFT) and compares Direct Preference Optimization (DPO), Proximal Policy Optimization (PPO) and Group Relative Policy Optimization (GRPO).

• DeepSeek-V3 describes a multi-token training objective, improvements to MoE load balancing, SFT, RL for fine-tuning, and framework optimizations related to training efficiency.

• DeepSeek-R1 discusses training techniques (large-scale RL, SFT, and RFT) for training R1-Zero and R1.

How DeepSeek R1-Zero was trained

The researchers apply RL directly to the base model (DeepSeek v3) without any cold start data to generate DeepSeek R1-Zero. However, the researchers note poor formatting/readability and the mixing multiple languages in the reasoning outputs.

In the training process, the researchers reused the Group Relative Policy Optimization (GRPO) algorithm which they introduced in their DeepSeekMath paper. This algorithm requires less computational overhead since it does not require a critic, and is stable as it measures output performance relative to the group. Responses received rewards based on accuracy and correct formatting.

R1-Zero was required to adhere to the format <think></think><answer></answer> when responding.

How DeepSeek R1 was trained

To address the usability issues in R1 Zero, the researchers expanded on R1 Zero with a four step process:

1. Supervised fine-tuning with a small number of R1-Zero outputs, cleaned by human annotators
2. Large scale RL, in the same way as R1 Zero. The language mixing issue returned, but the researchers mitigated it with a language consistency reward during training. This reward did result in a slight drop in performance, however.
3. Supervised fine-tuning data collected via rejection sampling. These prompts cover both reasoning and non-reasoning domains.
   • For reasoning domains (600k examples), only the correct response was retained.
   • For non-reasoning domains (200k samples) like writing, factual QA, translation, etc. the researchers reused techniques from DeepSeek v3: synthetic responses cleaned up with human annotators.
4. Further RL refinement for “helpfullness and harmlessness”:
   • For reasoning tasks, the researchers used rule-based rewards, as in R1-Zero.
   • For non-reasoning tasks, the researchers used reward models to capture preferences.
   • To optimize helpfulness, only the final summary was evaluated.
   • To optimize harmlessness, both the reasoning tokens and final summary were evaluated.

Further Distillation

Using the 800k inputs from step 3 above, the researchers generated outputs from R1. They then performed supervised fine-tuning on Qwen (1.5B, 7B, 14B, 32B) and Llama (8B and 70B).

Our elementary school models of the solar system really undersell how big space is. The problem is, space is too big and human brains are bad at exponentials. Logarithmic charts like this one are technically accurate, but my brain has a hard time contextualizing it.

To get an idea of how big space really is, let’s imagine that the earth is one millimeter wide. At that scale:

• The moon would be 3 cm away
• The speed of light would be about 2.36 cm per second
• The sun would be 10.9 cm wide (about the length of a toilet paper tube)
• The sun would be 11.7 meters away (as long as a school bus).
  • Note that the sun is only about 40% denser than water, so at this scale it would weigh only 0.689 kg (1.5 lb)
• Our farthest space probe, Voyager 1, would be 1.91 km (1.19 miles) away. It was launched in 1977, so at this scale it is averaging about 11 cm per day or 40 m per year.
• Proxima Centauri, the closest star, would be 3,162 km away.
  • For context, here is a time lapse of a 3,900 km flight between Los Angeles and New York City.

If aliens haven’t visited Earth yet, the distance between us and even the closest star is probably a good reason why.

class MyDataset(Dataset):
  
  # ... boilerplate ...

  def __getitem__(self, idx):
    item = self.data[idx]
    return item['anchor'], item['positive'], item['negative']

your collate function should be

def collate_fn(data):
    anchors, pos, neg = zip(*data)
    anchors = tokenizer(anchors, return_tensors="pt", padding=True)
    pos = tokenizer(pos, return_tensors="pt", padding=True)
    neg = tokenizer(neg, return_tensors="pt", padding=True)
    return anchors, pos, neg

and you can use it like

dataset = MyDataset()
dataloader = DataLoader(dataset, 
                        batch_size=4, 
                        shuffle=True,
                        pin_memory=True,
                        collate_fn=collate_fn)

for anchors, positives, negatives in dataloader:
  anchors = anchors.to(device)
  positives = positives.to(device)
  negatives = negatives.to(device)
  
  # do more stuff

How does the collate_fn work?

Under the hood, your DataLoader is making multiple calls to dataset.__getitem__(), one for each item in batch_size. Usually it is smart enough to know how to combine them into a tensor for training (here is the default collate function), but sometimes you need custom logic.

Technically you don’t need a collate function. You could collate within the training loop. But besides cluttering up your training loop, you won’t be able to take advantage of multithreaded processing (num_workers > 1), so your collation logic will block your training code.

What about pin_memory?

The pin_memory parameter specifies whether we should create the tensors in pinned RAM or not. This avoids some additional overhead when copying data to the GPU.

A common misconception is that this is pinning the GPU VRAM. Actually, this pins the RAM, not the VRAM.

If we check the Nvidia technical blog they give the following four recommendations (summarized):

• Minimize back and forth data transfers between RAM and GPU. Even if code is no faster on the GPU, it may be worth executing there to avoid the round-trip transfer cost, which is relatively slow.
• Use page-locked (“pinned”) memory when possible. By default, your machine uses pageable memory which means the physical address of RAM can change. When transferring data to the GPU, CUDA first needs to copy it to a fixed “pinned” address to be copied over. In your DataLoader, if you specify pin_memory=True, the tensors will automatically be created in pinned memory, avoiding the pageable memory -> pinned memory copy.
• Combine multiple small transfers into a large transfer to avoid overhead.
• Data transfers can be executed asynchronously. In PyTorch that looks like tensor.to(device, non_blocking=False).

In God we trust. All others must bring data. ~ W. Edwards Deming

Datasets that fit in memory

For simple machine learning problems, your PyTorch dataset class probably looks something like this:

class SimpleDataset(Dataset):
    def __init__(self, features, targets):
        self.features = []
        for feature in features:
            self.features.append(self._feature_transform(feature))
        self.targets = targets

    def _feature_transform(self, feature):
        # Optional feature transformation function which 
        # converts each feature into its input representation 
        # for the model. This might be an expensive operation, 
        # so its best to do now rather than during training.
        return some_transformation_fn(feature)

    def __len__(self):
        return len(self.features)

    def __getitem__(self, idx):
        return self.features[idx], self.targets[idx]

With this method, we basically load all of the data into RAM at once, which is perfectly fine for small datasets. But sooner or later you’re going to run into a machine learning problem with a large dataset. What do I mean by this? I mean a dataset which can’t easily fit into RAM/VRAM.

Datasets that fit on your hard drive

There are ways around this issue. One is using memory mapping. The idea is, rather than loading all of the data into memory, let’s store it on the disk and seek to it during training. This also means we can store the transformed tensors on the disk and load them as needed, rather than recomputing them every time we load data from disk.

Now we have a two step approach. The first step is transforming the features and targets to a memory-mapped file, and the second is loading it on-demand. For the first step we will create a list of memmap files, and for the second step we will use dask to virtually concatenate them together and read them as if they were a contiguous array.

import numpy as np
import uuid
import os

# store all features in memmap files
# you can do the same thing with targets

features_dir = 'features/'

for feature_file in feature_files:
    with open(feature_file) as f:
        features = transform_features(f.read())
    
    fname = f'{uuid.uuid4()}.memmap'

    # create the file on disk
    memmap_file = np.memmap(os.path.join(features_dir, fname), 
                            dtype='float64', 
                            mode='w+', 
                            shape=features.shape)

    # write features to memmap array
    memmap_file[:,:] = features

    # write changes to disk
    memmap_file.flush()
    

# Load the data from memmaps. We can use dask to create a 
# virtual contiguous array for simplicity.

import dask.array as da

class MemmapDataset(Dataset):
    def __init__(self, 
            feature_files, 
            target_files, 
            features_shape, 
            targets_shape):
        self.features = []
        self.targets = []
        
        for file in feature_files:
            self.features.append(np.memmap(file, 
                                           dtype='float64', 
                                           mode='r', 
                                           shape=features_shape))
        for file in target_files:
            self.targets.append(np.memmap(file, 
                                          dtype='float64', 
                                          mode='r', 
                                          shape=targets_shape))

        self.features = da.concatenate(self.features, axis=0)
        self.targets = da.concatenate(self.targets, axis=0)

    def __len__(self):
        return self.features.shape[0]

    def __getitem__(self, idx):
        return (
            torch.tensor(self.features[idx]), 
            torch.tensor(self.targets[idx])
        )

Sparse vectors

For very sparse vectors, where most of the items are zero, my recommendation is to store the nonzero indices of each vector rather than the whole thing. This is obviously an empirical determination, but it can require significantly less space to store the indices and build at train time than to store the entire vector. You can use a custom collate function to expand the training data during training:

def collate_fn(data):
    '''
    Here, data is a list of tuples returned from dataset[idx] 
    above. But supposing that "features" actually corresponds 
    to the nonzero indices of the features, we can expand the 
    tensor for training.

    This example is for a model which expects input dimension
    (bs, 1024) but you can adapt for your needs.
    '''
    feature_idxs, targets = zip(*data)
    bs = len(features)
    features = torch.zeros((bs * 1024))
    features[feature_idxs] = 1.0
    return features.view((bs, 1024)), targets

You then use the collate function like

dataset = MemmapDataset(...) # or SimpleDataset
dataloader = DataLoader(dataset, 
                        collate_fn=collate_fn, 
                        batch_size=bs)

Datasets that don’t fit on your hard drive

Sometimes the dataset is so big that it doesn’t fit on your hard drive. In that case, you might need to reorganize your code a bit, preprocessing your data into unique UUID-based shard files (sids), and training each sid one by one.

sids = [
    # you may have thousands of shard ids here
    '5c91ad9e-4963-4dfa-8885-6a351dd9fbb8',
    '8ecb542f-00df-468e-8f89-34a6608648d6',
    'f4536d05-a8ff-4591-9322-10860cd06942'
]

for sid in sids:
    dataset = ShardedDataset(sid)
    dataloader = DataLoader(dataset, batch_size=64, shuffle=True)

    for i, batch in enumerate(dataloader):
        # do your training
        pass

And your Dataset class will look something like

class VeryLargeDataset(Dataset):
    def __init__(self, sid):
        self.data = self._load_shard(sid)

    def _load_shard(self, sid):
        # load from your hard drive or remote
        pass

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]

Despite the reorganization required, this is a very flexible approach and can scale to an effectively infinite amount of data. For continuous training, this is often a very effective approach.

Bonus: Use Rust/C Python bindings

Your feature transformation might be pretty slow depending on what you need to do. In my experience I was waiting 10 minutes to process text files in Python, which really slows down the development/training process. A much more effective approach can be to use Python bindings from a faster language, such as Rust or C.

Rust can even return native numpy arrays, making data transformations very efficient. For more on setting up Rust bindings, check out my how-to article.

import my_bindings

files_to_process = [...]
data = []
for file in files_to_process:
    data.append(my_bindings.process_file(file))

I will demonstrate rust bindings for summing the integers in a large text file containing a billion digits that looks like

6,9,8,3,0,1,8,4,9,7,6,3,4,2,6,0,0,5,1,1, . . . ,4,5,9,3,3,2,8,3

General steps

1. install rust and maturin
2. set up boilerplate
3. add your function
4. compile and import

Install Rust and Maturin

curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
pip install maturin

See also install rust and install maturin

Set up boilerplate

# select "pyo3" here
maturin new cool_rust_bindings
cd cool_rust_bindings

Add your function

Open src/lib.rs and change it to the following:

use pyo3::prelude::*;
use std::fs::File;
use std::io::{self, Read};

#[pyfunction]
fn sum_file(filename:&str) -> io::Result<i64> {
    let mut file = File::open(filename)?;
    let mut contents = String::new();
    file.read_to_string(&mut contents)?;

    let mut total:i64 = 0;
    for x in contents.chars() {
        if x != ','{
            total += (x as i64) - 48;
        }
    }

    Ok(total)
}

#[pymodule]
fn cool_rust_bindings(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(sum_file, m)?)?;
    Ok(())
}

Here we’re just creating a simple function that sums the integers in a file. The rust code gave a 3x speedup over the corresponding python code. Depending on the task, I’ve gotten 10-20x speedups.

Compile and import

Run the following from the cool_rust_bindings directory

maturin develop

This will add cool_rust_bindings to your python path. Use it like

import cool_rust_bindings

cool_rust_bindings.sum_file('large_file.csv')

I spent some time considering words to describe a chess position with many legal moves. “Complex” doesn’t quite capture the situation since we would usually describe a complex position as one with many tactical interations. Ramified seems to make the most sense, as it describes “branching out”.

The opening position in chess has 20 legal moves. From there, the number of legal moves in a position tends to increase as pieces move towards the center of the board, before decreasing as the number of pieces on the board drops in the endgame.

People have composed artificial positions with a huge number of legal moves. The most ramified artificial position I have found has 218 legal moves. But it’s hard to imagine how such a position could come about in a real game, which is what I’m interested in. So I churned through all of the positions in all of the Lichess games which took place in 2023 in search of the most ramified chess position. Here is what I found.

Most games max out around 16 moves in

In the chart below you can see the maximum moves at any point in the game. The x-axis represents half moves. (In other words, if white makes a move and black makes a move, that’s two half moves.) Each game starts out with exactly 20 possible moves for white, which quickly increases in the middle game before gradually decreasing again.

The colors represent the percentiles, since at every slice along the x-axis there is actually a distribution of positions and associated legal moves. Interestingly, we can see that for the 10th percentile the bottom really falls out after 40 moves or so, with positions so constrained that there are only a handful of legal moves. These likely represent positions where the king is in check.

100+ move positions are rare

To parse the games, I wrote a multithreaded rust script with python bindings to process chunks of 5000 games at a time. Within a chunk of games, it was pretty typical for the maximum number of legal moves in any position in any game in the chunk to be between 70 and 100.

Here is a typical position with 84 legal moves for white:

Very few ever exceed 120

When you get past 100, most of the positions have several queens. There are exceptions, however. This 112 move gem has zero.

Here is a typical 120-move position with six queens:

The 147 move position 🤯

Here it is, the 147 legal move masterpiece.

Runners up

• 8 queens, 146 legal moves
• 9 queens, 144 legal moves
• 9 queens, 142 legal moves
• 8 queens, 142 legal moves

TLDR: A sparse autoencoder is just a regular autoencoder that encourages sparsity with an L1 penalty or KL divergence loss rather than using a low-dimensional bottleneck.

If you understood all of those words above, you may be interested in the OpenAI paper which used sparse autoencoders to interpret features from GPT-4.

If not, I’ll try to break it down.

What is an autoencoder?

An autoencoder is a machine learning architecture which contains two functions: an encoder and a decoder. The encoder learns to create an efficient representation of the input, similar to lossy file compression. The decoder learns to reconstruct the original input from the efficient representation.

I have a couple of other pretty good explanations on this. In the first post, I described how to build an autoencoder. Later, I discussed their potential applications in vector retrieval.

What is sparsity?

Sparsity just means that most of the values in a vector are zero.

What is an L1 penalty?

L1 stands for penalizing the $ L^1 $ norm of a vector. In other words, the sum of the absolute values of the values in the vector:

$$ \text{L1 norm} = \sum_{i=1}^N \vert x_i \vert $$

What is KL divergence?

KL divergence is a measure of the “distance” one probability distribution has from another:

$$ \sum_{x \in X} P(x) \log \left( \frac{P(x)}{Q(x)} \right) $$

Note that KL divergence is not symmetrical, so distance isn’t an entirely accurate description. If you’re looking for a symmetrical measure of distribution similarity, use Jensen-Shannon divergence. Additionally, $ X $ doesn’t have to be a continuous variable, although it often is.

In the demo below you can see that as the probability distributions overlap more and more, the KL divergence decreases.

What is a KL divergence loss?

KL divergence loss tries to force a model’s predicted distribution to match a target distribution. Common applications of KL divergence loss are t-SNE dimensionality reduction and generative models like variational autoencoders and generative adversarial networks.

Naive and semi-naive approaches

One approach might be to compare every vector and take the argmax. In that case, for vectors of length D our runtime will be 1 billion x D.

def search(needle, haystack, arg_k=1):
    # normalize needle and haystack vectors
    needle /= np.linalg.norm(needle)
    haystack /= np.linalg.norm(haystack, axis=1, keepdims=True)

    # compute cosine similarities
    similarities = np.dot(haystack, needle)

    # get the indexes of the top k elements
    idx_topk = np.argpartition(a, -k)[-k:]

    # return haystack vector with highest cosine similarity
    return haystack[idx_topk]

In practice, there are some optimizations we can try which can increase the practical speed of this operation:

• We can parallelize computations on-device with SIMD or BLAS.
• We can parallelize computations on-device with GPU.
• We can parallelize across cores for machines with multiple cores.
• We can parallelize across devices. For example, partitioning the database across 10 clusters, finding a “winner” from each cluster, and then repeating the search across the 10 winners.

There are also some relatively simple but “destructive” techniques we can also try:

• We can try pruning the dataset in some way beforehand. For example, maybe we don’t need to consider old or unpopular vectors.
• We can try reducing the dimensionality of the vectors with something like PCA.
• We can try quantizing the vectors e.g. f64 -> f16 reduces memory requirements 4x. (Usually this requires quantization-aware training.)

Hierarchical Navigable Small Worlds (HNSW)

Sometimes linear time is still too slow. In that case, we can consider an approximate solution, where some acceptable percent of vectors returned are actually not in the top K nearest neighbors. This is called approximate nearest neighbors (ANN). If we have a large system, it’s possible to use another more precise model to rerank the vectors results anyways. In that case, we just need to set K large enough to reliably include the vectors we really need.

ANN algorithms typically come in three flavors: graph based approaches, space partitioning based approaches, and hash based approaches. If your dataset is large enough, you might add compression (e.g. product quantization) on top of it.

Intuition

HNSW is a bit of a Six Degrees of Kevin Bacon approach. I’ll give what I think this is a fairly good intuition for how it works.

Suppose you have 1,000 friends and you’d like to know which five of them live closest to a landmark like the Golden Gate Bridge. You don’t have everyone’s exact address, but each friend keeps a list of their close friends’ addresses. You might start with a friend who lives in the US. From there, you check their friends to find who lives closest to the target, and then search among their friends for someone even closer. By using a hierarchical approach, you don’t need to search through everyone, just those likely to be closest.

In the analogy, the Golden Gate Bridge is the “needle” and the friends are the “haystack”. This is a pretty efficient method, but of course it is possible to miss a node if their friend wasn’t linked to an earlier friend.

Searching

Searching is fairly straightforward. Start from an entry point (maybe a random point) in the top layer. For each of the neighbors of the entry point, compute the distance between the query and that point. If any of the neighbors are closer to the query, hop to that neighbor and repeat the process. Select the closest point and go to the next layer down. We keep searching until we reach layer 0.

To compute the K nearest neighbors during search (with K > 1) we simply maintain a min heap, and add all nodes considered during the search. The number of nodes returned in this search (K) is also known as ef_search in the HNSW paper.

Addition

To add a vector X, we start by computing the highest layer L the node will appear in. L is randomly selected using an exponentially decaying probability. (All nodes appear in level 0, but can have skip links to higher levels.) Next, starting from the top layer, we traverse the graph structure as if we were searching for X.

Once we reach layer L from earlier, we begin to make connections. Connections can be made with previously found nearest neighbors or nearest neighbors on the current layer. The number of nearest neighbors to consider connections with is controlled by the ef_construction parameter. The number of links to actually create is controlled by the parameter M.

Performance considerations

• HNSW can require a large amount of memory. This can be reduced with product quantization, a lossy vector compression method.
• Adding new elements is slow. There is a tradeoff between index quality and time required to build the index controlled by ef_construction.

More information about HNSW parameters can be found here.

Embeddings, or vector representations of a document (which could be a piece of text, image, sound, etc.), can be extremely useful for making sense of large datasets. They transform information into a vector space such that their distance corresponds to their similarity.

Enterprising readers might be asking themselves how to get these vectors (also known as embeddings) in the first place. One way is to simply pay for them. This isn’t ideal for a couple of reasons:

1. It can be expensive. You’ll need to pay once to embed each document, and separately for each query. If you have a large corpus, it can be cost-prohibitive. At the time of writing, OpenAI charges about $130 per million tokens (around 1000 paragraphs) for their largest model.
2. “Similarity” may mean different things depending on your use case. For example, suppose we are retrieving documents for customer support. For the embedding model to learn that two documents share user behavior characteristics (i.e. two documents were opened by support agents in the same session), that information needs to be available in the training process.

Below is an overview of three of the main training regimes I have used for creating embeddings. For more information and in-depth examples, I highly recommend the loss overview page of the Sentence Transformers library. Generally speaking, there are three categories of methods: unsupervised methods, contrastive learning methods (positive/negative labels), and regression methods (floating-point labels).

Autoencoders

One approach for vectorizing a document, image, or other blob of information is to simply use an autoencoder. An autoencoder is a function which learns a lossy compression function. It can be considered an unsupervised method since each item is its own label.

However, although they are conceptually pretty simple, autoencoders aren’t always best. They may learn an efficient representation of your document, but it may be wasting a lot of space on things you don’t care about.

For example, if it was reproducing a picture which contained a TV showing static, you may not care exactly what the static looks like, just that it has static. So it’s probably a waste of space to try to reproduce every pixel of static. Likewise, if it were encoding people, it might be really important to you that the people have 5 fingers. Unfortunately, the autoencoder wasted too much space encoding the TV and not on boring details like numbers of fingers.

See also: the “noisy TV problem”

Fine-tuning with Sentence Transformers

The applicability of the next few methods will depend on the format of your data and labels.

Here is a simple example for creating text embeddings with sentence-transformers using ContrastiveLoss:

from sentence_transformers import (
  SentenceTransformer,
  SentenceTransformerTrainer
)
from sentence_transformers.losses import ContrastiveLoss
from datasets import Dataset

model = SentenceTransformer('all-MiniLM-L6-v2')
train_dataset = Dataset.from_dict({
    "sentence1": [
      "I feel the need...",
      "Life is like a box of chocolates."
    ],
    "sentence2": [
      "...the need for speed!", # similar
      "Here's Johnny!" # dissimilar
    ],
    "label": [1, 0],
})

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=ContrastiveLoss(model),
)
trainer.train()

And to generate the embedding:

embedding = model.encode(["There's no crying in baseball!"])
print(embedding.shape) # (1, 384)

Contrastive Learning: Triplet loss

Triplet loss explicitly learns embeddings to be used in cosine similarity or euclidean distance (L2 norm) comparison. It uses triplets of the form (A, P, N), where A is an anchor, P is a positive example which is similar to A, and N is a negative example which is dissimilar to A. For example if we wanted to learn embeddings for songs, we might say a positive example is a spectrogram clip from the same artist as the anchor, and a negative sample is from a different artist.

$$ \text{L} = \sum_{i=1}^N \text{max} \left(0, cos( f(a_i) - f(p_i) ) - cos( f(a_i) - f( n_i )) + m \right) $$

for cosine similarity or for euclidean distiance,

$$ \text{L} = \sum_{i=1}^N \text{max} \left(0, \lVert f(a_i) - f(p_i) \rVert_2 - \lVert f(a_i) - f( n_i ) \rVert_2 + m \right) $$

In other words, we enforce that the distance between the anchor and the negative sample is greater than the distance between the anchor and the positive sample plus some margin. For cosine similarity, what this would look like is that your vectors appear as points on the perimeter of a high-dimensional clock face (i.e. the surface of an N-dimensional hypersphere). Positive samples grow closer to the anchor while the negative samples are pushed away. The higher the margin, the closer positive samples will be pushed together.

Contrastive learning: Contrastive loss

If you have binary labels, you can use contrastive loss. It works with positive pairs (anchor, positive) and negative pairs (anchor, negative), denoted with a binary label. In the original paper, $ y \in [1,0] $ is the label indicating the desired distance between the points (not whether they are similar).

Like triplet loss, contrastive loss also includes a margin parameter. For dissimilar pairs, a loss is incurred only if their predicted distance falls within the margin’s radius.

$$ \text{loss} = \frac{y_{i}(D_{W_i})^2 + (1 - y_i) \left( \text{max}(0, m - D_{W_i}) \right) ^ 2 }{2} $$

Where $ D_{W_i} $ is euclidean distance between the two embeddings of sentences $ a_i $ and $ b_i $:

$$ D_{W_i} = \lVert f(a_i) - f(b_i) \rVert _{2} $$

Note that triplet loss is often confused with contrastive loss, but they are not the same.

Regression: Cosine similarity loss

If you already have some numerical labels for the similarity of two pieces of text, you can compute mean squared error between their predicted cosine similarity and the actual similarity.

For pairs of sentences a_i and b_i and label y_i,

$$ \text{loss} = \frac{1}{N} \sum_{i=1}^{N} \left[ cos\left(f(a_{i}), f(b_{i})\right) - y_{i} \right] ^ 2 $$

Regression: CoSENT loss

CoSENT loss is an improved version of cosine similarity learning which makes better use of in-batch values. For a batch with size N it makes up to Nx(N-1)/2 predictions. (I.e. for a batch size of 10 we make up to 45 predictions rather than the 10 in Cosine Similarity Loss.) It is essentially a ranking loss, taking advantage of the knowledge of the relative similarities within the batch.

$$ \text{loss} = \log \left( 1 + \sum_{\text{sim}(i,j) > \text{sim}(k,l)} e^{\lambda (\cos(k, l) - \cos(i, j))} \right) $$

Experiments have shown improved performance over Cosine Similarity loss.

Using them: Sub-linear Vector Search

Once we have generated our vectors, we will want to search through them as quickly as possible. If our application is maintaining a large (i.e. million or billions of items) catalogue, linear scan through them all will be too slow.

Instead, we can use an approximate nearest neighbors algorithm like HNSW.

🌟Check out the demo (in two dimensions) below. Centroids are colored white.🌟

Note that the points are changing color only, not moving.

General algorithm

The basic K-means algorithm is fairly simple and has two steps, repeated until convergence (i.e. when no points change cluster):

1. assign points to the cluster corresponding to closest centroid
2. update the centroid locations to the mean of all points assigned to the associated cluster

import numpy as np

'''
General setup: we will demo 5000 points, each with two 
dimensions. They will be clustered into 5 clusters.
'''

N, k, d = 5000, 5, 2
points = np.random.random((N,d))
centroids = np.random.random((k,d))

while True:

  # assign each point to closest cluster
  clusters = assign_clusters(points, centroids)

  # recalculate centroids
  new_centroids = update_centroids(points, clusters, k)

  # if centroids haven't changed, we're done
  if centroids_equal(new_centroids, centroids):
    break

  # otherwise, update centroids
  centroids = new_centroids

Vectorizing the algorithm

Now that we’ve established the general algorithm, we might consider how it can be vectorized. In essence, this means we’ll avoid writing for-loops in python. In numpy, this means cython will handle for-loops. For environments with a GPU, we can use PyTorch, JAX or Tensorflow to take advantage of parallel processing that GPUs excel at.

How broadcasting works

Before going further it would be good to pause and introduce broadcasting. The numpy guide is much more comprehensive, but the gist is that because numpy needs tensors to have the same shape before performing certain arithmetic operations on them, numpy will automatically try to align their shapes. Importantly, this is done without making copies of the data. The steps for broadcasting are as follows:

0. A scalar will simply take the shape of the vector it’s being broadcasted to.
1. If one vector has fewer dimensions, left-pad its shape with 1s.
2. Starting from the right, check dimension compatibility. If they are equal or one of them is 1, the dimensions are compatible.

For example,

# scalars are the simplest
(np.ones((2,2)) * 3).shape # (2,2)

# (3,) is left padded to (1,3)
(np.ones((3)) * np.ones((2,3))).shape # (2,3)

# if dimensions equal or one of them is 1, they are compatible
(np.ones((2,3,1)) * np.ones((2,1,4))).shape # (2,3,4)

# error! [1,1] * [1,1,1] doesn't work
(np.ones((2)) * np.ones(3)) # incompatible shapes!

For more information, check out the Numpy guide: https://numpy.org/devdocs/user/basics.broadcasting.html.

assign_clusters()

Ok, back to K-means. Let’s start with assign_clusters(). A point is assigned to the cluster corresponding to the centroid which is closest.

The general idea is that for each point we’d like to know its distance to each centroid. We can then take the argmin of those distances to identify the correct cluster and return a result with shape (N,) (N = # of points). Distance will be computed using euclidean distance i.e. subtract, square, and sum.

First we expand the dimensions of the points from (N,d) to (N,1,d). We also expand the dimensions of the centroids to (1,k,d). (This happens automatically via broadcasting.) When we take the difference, we are left with a vector with shape (N,k,d). For each point in N, this yields a vector with k rows and d columns.

Now that we have the differences, we square them, and take the sum across the last dimension. This yields a vector with shape (N,k).

For actual euclidean distance we would then take the square root of the result, but since we just care about the relative distances we can skip this step.

Finally, we just take the argmin of the last dimension, i.e. for each row, the index of the smallest distance. Now we have our (N,) result.

def assign_clusters(points, centroids):
  # (5000,2) -> (5000,1,2)
  points = np.expand_dims(points, 1)

  # (5000,1,2) - (5,2)
  # with broadcasting becomes (5000,5,2)
  diff = points - centroids

  diff = diff ** 2

  # (5000,5,2) -> (5000,5)
  diff = diff.sum(-1)

  # (5000,5) -> (5000,)
  diff = diff.argmin(-1)

  return diff

update_centroids()

To update the centroids, we simply take the mean of the points within that cluster. To do this, we will create a vector with shape (N,k,d) again, and mask out points not in the cluster.

We start by creating an N x k mask, where 1s correspond to the n-th point being in the k-th cluster. We will also temporarily add a dimension at the end to make this N x k x 1.

Adding a dimension to the points vector for proper broadcasting, we can then create the N x k x d vector we planned above. Summing this across the first dimension yields a k x d vector, representing the sums of all of the points within each of the k clusters.

Now we just need to divide by the number of points in each cluster, which we can get by summing the N x k mask we created earlier across the first dimension, yielding k counts.

At this point, we can divide the sum by the counts to get the means.

def update_centroids(points, clusters, k):
    N = clusters.shape[0]

    # create an N x k mask, where each row contains a 1
    # in the position corresponding to its cluster
    mask = np.zeros((N, k))
    mask[np.arange(N), clusters] = 1

    # add a dimension in the middle so that 
    # multiplication will broadcast
    # (5000,2) -> (5000,1,2)
    points = np.expand_dims(points, 1)

    # (5000,5,1) * (5000,1,2) results in a (5000,5,2) vector
    # for each of the 5000 points, we have a 5x2 vector where only
    # one of the five rows will have non-zero values
    masked_points = np.expand_dims(mask, -1) * points

    # sum across points
    # (5000, 5, 2) -> (5, 2)
    sums = masked_points.sum(0)

    # sum across points
    counts = mask.sum(0)

    # avoid divide by zero
    counts[counts == 0] = 1

    # divide sum by count to get means
    return sums / np.expand_dims(counts, -1)

centroids_equal()

As a last implementation, we need to check whether the centroids have changed. For this we can just call np.array_equal() on the old and new centroids vectors.

def centroids_equal(vec1, vec2):
  return np.array_equal(vec1, vec2)

Comparing vectorization vs loops

To compare the vectorized implementation with the non-vectorized implementation, I ran each 100 times and took the mean clock time to converge. Note that because the points are random, K-means may take different numbers of iterations to converge.

The vectorized implementation (green) runs in significantly less time, between 7 and 21 times faster.

Learning to rank is an important problem in many machine-learning products such as search, recommendation, and advertising. Originally, many machine learning systems used simple logistic regression models, but it quickly became apparent that combining two or more features together was even better. This is called feature crossing.

A lot of research and engineering work has gone into learning useful feature crosses. The fundamental problem is that although higher-order feature crosses can be more informative, they are also more sparse, and the number of high order features grows exponentially. Some attempts to address this have been:

DCNv2

DCNv2 improves on the original DCN architecture in a few important ways. First, it uses a more expressive feature cross formula. Second, it introduces matrix factorization to dramatically reduce the required number of parameters in the model. Finally, the authors also introduce a gated mixture of experts setup to take advantage of the parameter savings.

Embedding layer

To create the embedding layer, all features are concatenated together to create a tensor with shape batch size x embedding size. For one-hot encoded categorical features, they can be projected to a much lower dimensionality with a learned matrix (i.e. a single dense layer with no bias term).

One convenient property of this embedding scheme is that feature order doesn’t matter. All features will be crossed with all other features, so we don’t need to worry about the way that features are added and removed.

Cross Net

The formula for a single cross operation is

"""
Dimension key:
    B: batch size
    D: embedding dimension

Params:
    x_0: original embedding, shape BD
    x_i: i-th embedding, shape BD
    weight_i: i-th learned weight matrix, shape DD
    bias_i: i-th learned bias term, shape D

Returns:
    The residual embedding after i+1 cross operations
"""

def cross(x_0, x_i, weight_i, bias_i):
    x_0 = x_0.unsqueeze(-1)         # (B, D, 1)
    x_i = x_i.unsqueeze(-1)         # (B, D, 1)
    bias_i = bias_i.unsqueeze(-1)   # (D, 1)

    kernel = weight_i @ x_i     # (D, D) @ (B, D, 1) = (B, D, 1)
    kernel = kernel + bias_i    # (B, D, 1) + (D, 1) = (B, D, 1)
    
    result = x_0 * kernel + x_i # (B, D, 1)
    return result.squeeze(-1)   # (B, D)

Note that the weight_i @ x_i + bias_i term is equivalent to a fully-connected neural network layer with no activation. To this result, we element-wise multiply the original embedding, then add the i-th embedding.

Low rank approximation

Of course, with high embedding dimensionality, the weight matrix can easily grow to millions of parameters, and we need a weight matrix for each of the cross operations. A common technique to reduce the number of parameters is to use a low rank approximation of the weight matrix. In other words, we can replace the weight @ x operation with weight_u @ (weight_v @ x)

In particular, rather than requiring DxD parameters for each cross operation, we will now only need DxR + RxD = 2xRxD parameters, where R is the low-rank approximation dimension. So if R is less than half of D, we can save parameters using this technique.

"""
Dimension key:
    B: batch size
    D: embedding dimension
    R: low rank dimension

Params:
    x0: original embedding, shape BD
    xi: i-th embedding, shape BD
    weight_u_i: i-th learned weight matrix, shape DR
    weight_v_i: i-th learned weight matrix, shape RD
    bias_i: i-th learned bias term, shape D

Returns:
    The residual embedding after i+1 cross operations
"""

def low_rank_cross(x_0, x_i, weight_u_i, weight_v_i, bias_i):
    x_0 = x_0.unsqueeze(-1)         # (B, D, 1)
    x_i = x_i.unsqueeze(-1)         # (B, D, 1)
    bias_i = bias_i.unsqueeze(-1)   # (D, 1)

    kernel = weight_v_i @ x_i       # (R, D) @ (B, D, 1) = (B, R, 1)
    kernel = weight_u_i @ kernel    # (D, R) @ (B, R, D) = (B, D, 1)
    kernel = kernel + bias_i        # (B, D, 1) + (D, 1) = (B, D, 1)
    
    result = x_0 * kernel + x_i     # (B, D, 1)
    return result.squeeze(-1)       # (B, D)

Mixture of low-rank experts

The parameter savings from using low-rank approximations of the weight matrices can be reinvested into additional parameters. In particular, if we have a gating function G(x) we can use it to decide between a group of “experts” which are themselves low-rank Cross networks.

Note that setting G(x)=1 and K=1 reduces the mixture of experts to the original low-rank DCN architecture.

Analysis of hyperparameters

The number of cross layers (i.e. cross depth), the choice of low-rank approximation, and the number of experts can each have a large effect on the model performance. For the depth of the crosses, the researchers found that performance improved with the number of cross layers, but diminishing returns as depth increased.

The behavior of the matrix rank was interesting. For very low rank, the model performed similarly with other baselines. However, increasing from rank 4 to 64 showed dramatic performance improvement, and notable but lesser improvement afterwards. This suggests there is some “elbow” (the authors call this a threshold rank) which maintains good performance without consuming too many parameters.

For the number of experts, the researchers found that the total rank sum of all of the experts was the main contributing factor for performance. In other words, 1x256, 4x64, 8x32, 16x16 and 32x8 all performed similarly. They hypothesized that their simple gating function was responsible for this.

Results on synthetic datasets

To better understand how the DCN, DCNv2 and DNN are able to model cross interactions, the authors created synthetic functions for the various architectures to learn. They created three classes of functions with increasing difficulty to learn.

Note that the third equation can be modeled by a single cross operation.

In their experiments, the authors found:

• The DCN consistently used the fewest parameters (maximum of 300 parameters) and was able to accurately model the simpler functions, but was unable to accurately model the third, most complicated function with a single layer. However, since a single layer was only 300 parameters, experimenting with a larger model would have been interesting.
• The DNN used the most parameters (up to 758k parameters) and was unable to model any of the functions as accurately as either DCN or DCNv2.
• The DCNv2 used up to 10k parameters for the most complex function and achieved the lowest RMSE loss by far.

The authors also directly compared DCNv2 models with varying numbers of layers with a traditional DNN in modeling a more complex function containing sine terms and polynomials of order 1-4 and found that DNNs are highly inefficient in modeling higher order interactions.

Results on Criteo

Criteo is a popular click-through-rate (CTR) prediction dataset containing 7 days of user logs. It contains 13 continuous features and 26 categorical features.

On Criteo, DCN-v2 (best setting using two layers) had a 0.21% better AUC than a DNN baseline. DCN-Mix (best setting using 3 layers, 4 experts, and low rank of 258) was 0.17% better on AUC than the DNN baseline.

Results on MovieLens-1M

MovieLens-1M is a popular dataset for recommendation systems research, containing one million samples of [user-features, movie-features, rating] triplets. The task was treated as a regression problem. Ratings of 1 and 2 were set to 0 and 4 and 5 were set to 1. 3s were removed.

On MovieLens, DCN-V2 was 0.24% better on AUC than the DNN baseline. DCN-Mix was even better, at +0.39% over the DNN baseline.

Qualitative answer

There are many qualitative answers to this question, especially when comparing “blunders” and “mistakes”:

• “a move that negatively affects their position in a significant way” ~ chess.com
• “severely worsens the player’s situation by allowing a loss of material, checkmate, or anything similar” ~ Wikipedia
• “Blunders tend to be immediately refutable, while mistakes require planning to capitalize on.” ~ r/chess

More generally, qualitative analysis is happy to consider more abstract aspects of a position than piece count, such as the positional strength of an active knight, king safety, and “practical chances” i.e. how easy it is for a human to play the position.

An issue with these qualitative answers is that while their words may be correct, smart people may still disagree with their applicability at the margins. For a suboptimal move to have a “significant” negative effect, it requires that the opponent notices and takes advantage of it.

Quantitative answer

The quantitative answer considers a move which causes a significant drop in probability of winning to be blunder. What is “significant”? A change 14% or greater.

How is winning probability in a chess game calculated? Objectively, since there are only three possible outcomes in a game (win, draw, loss), by definition any real advantage will lead to a win with perfect play. But objectively, humans aren’t perfect. Even grandmasters can let an advantage slip. If Magnus Carlsen doesn’t capitalize on your blunder, was it really a blunder?

From a machine learning perspective, we can view winning probability as a logistic regression problem, where centipawn evaluation is a feature, and game outcome is a label. If we further limit data points to 2300+ rated games, this is what Lichess uses.

However, Lichess’ accuracy article suggests that this relationship is slightly flawed in two ways:

1. Higher-rated players tend to play more theoretical openings which will have a lower centipawn loss.
2. Higher-rated players tend to resign in lost positions, perhaps underestimating their true win percentage due to opponent blunders. Never resign!

Of course, this isn’t perfect, and there’s an argument to be made that outcomes for 2300+ elo players may not be representative of lower-rated players. It also doesn’t take time pressure into consideration. Taken to the extreme, we can even use an adversarial player-modeling algorithm to win faster than Stockfish. But there is a tradeoff between accuracy of our metric and generalization power of the model.

(We can also go the opposite direction and directly try to quickly predict the stockfish evaluation of a position.)

One other important objection to this line of inquery is that the centipawn evaluation of a position is not a constant. The evaluation varies by search depth and between engines. So if Stockfish 15 evaluates a position at +1 and Stockfish 16 evaluates the same position at +1.5, white’s actual winning chances haven’t changed at all. The evaluation is not anchored to any real value, especially with the introduction of NNUE eval functions.

This variant of the Game of Life uses three overlapping channels, so instead of just one simulation, there are three simultaneous simulations. I visualize these in the three color channels, red, green and blue. Two or more channels active on the same cell are represented with additive color mixing.

These additional channels enable additional interaction between channels, which is pretty neat. I’ve added a rule which revives a dead cell if the sum of its active neighbors from other channels is 9. One nice result of this inter-channel interaction is that stagnant colors can be reactivated by other colors that pass by.

Finding parameters that don’t cause the simulation to blow up or die out is a bit of a trial and error. I made an empirical choice which results in a lively “pond” after the chaotic initial state. I’m sure there are other nice combinations as well.

Other variants I experimented with briefly were to include more than three channels and to consider more complex relationships between the channels. However these ran slower and/or were harder to balance between chaotic and desolate tendencies.

This was admittedly a fairly stupid experiment on the face of it. ChatGPT is a decoder-only model. It shouldn’t be able to perform an image recognition task. But then again, a decoder-only model wouldn’t have been my first choice for translation or summarization either. In my experience, ChatGPT has created translations which are at least as coherent and idiomatic as Google Translate, if not more so.

So I thought, why not give it a shot?

Process

MNIST is a computer vision dataset containing thousands of handwritten digits along with their actual labels. Images are represented as 28x28 tensors where each element of the tensor represents a intensity pixel intensity between 0 and 1. Traditionally this is formulated into a classification problem, with the goal of choosing the correct class out of the 10.

To simplify the input, I flattened the tensor and translated the pixel intensity values to the range 0-10. This allowed me to use fewer tokens at the expense of some of the gradations between light and dark.

I gathered 10 random examples from each of the 10 classes (100 total images) and fed them into gpt-3.5-turbo with the following prompt:

The following is a flattened representation of an image of a handwritten digit. The image was 28x28 but has been flattened to 1x784. Each number represents a pixel intensity from 0-10. Please tell me which digit the following pixel list represents:

Followed by the list of 784 pixel intensities.

Results

ChatGPT scored 11/100. Barely better than guessing. A summary of the sessions can be found at this gist. A visual summary is below, with the blue signifying correct guesses.

Note: although the image shows the digits sorted, they were prompted in random order.

It is interesting that only the 0s and the 7s were correctly identified. Let’s take a look at the distribution of ChatGPT’s guesses.

Oh. It got 8/10 zeros correct because it guessed zero 85% of the time.

Typically, antennas are designed and built by hand by domain experts. This is a very time-consuming process, however, so researchers have been investigating evolutionary algorithms since the 1990s. Inspired by natural evolution, an evolutionary algorithm is based on small, random changes and an evaluation metric. In this paper, the authors describe the use of an evolutionary algorithm to design an antenna for a small satellite weighing only 25 kilograms called ST5.

The researchers then describe the technical specifications which the antenna would need to satisfy. The antenna needed to weigh under 165 grams and have a height and diameter of around 15 centimeters.

In representing the antenna, the researchers used a tree structure, with each node capable of representing one of several operations: forward, rotate-x, rotate-y, and rotate-z. The “forward” operation adds a length of wire with a given radius. The “rotate-x” operation rotates the current state about the x-axis. The “rotate-y” and “rotate-z” operations do the same for the y and z-axes.

The researchers used a fitness function which was the product of VSWR, a gain_error term and a gain_smoothness term. The gain_error term is similar to the least squares error function. The gain_smoothness term describes how uniform the gain pattern was, since the satellite would be spinning. All three of these factors were multiplied together to calculate an overall fitness score against which all candidate antennas were measured.

Due to a change in the orbit of the satellite, the specifications for the antenna were updated and the researchers needed to modify their evolutionary algorithm. This was completed within one month. The evolved antenna consumed less power, took less time to build, was less complex, and performed better than traditionally designed antennas. Since the satellite had 2 antennas, the researchers measured the combined performance. The evolved antennas were 93% efficient while the designed antennas were only 38% efficient. Additionally, the evolved antennas were much faster to design and fabricate.

The researchers then designed another antenna for NASA’s TDRS-C satellite. This design used an evolutionary algorithm in combination with a stochastic hill-climbing algorithm. They used a loss function which attended to the standing wave ratio and gain at several frequencies relevant to the satellite. In performing the evolutionary algorithm, 150 algorithm processes were run for 50,000 iterations after randomizing many parameters. After this first stage, the best antenna from the previous 150 was optimized using stochastic hill climbing with random mutations. From this second stage, the 23 best antennas were selected and run through another stochastic hill climbing step for 100,000 iterations. Of the 23 finalists, one of the evolved antennas exceeded the project specifications and was optimized further using more accurate software.

Background

Search requires an automatic way to find the most relevant documents to a query. There are bag-of-word approaches to this task (for example BM25) and neural approaches. An example of a bag-of-words approach might simply be to count the number of similar words between the query and each document, and return the document with the highest number of similar words. There are word-stuffing issues with this idea, but the larger issue is that a bag-of-words strategy can’t account for synonyms. If I search for bad guy I will never find villain without some additional logic to account for this. A neural network implicitly understands the relationship between words, and avoids the fragile logic of simple word counts.

The idea of a neural encoding approach is pretty simple. For each document in your corpus, pass the query a document into a function which will return a similarity score between 0 and 1. Then, just sort the documents by that score. There are two main architectures for doing this: dual-encoders and cross-attention models.

The great thing about DE models is that document embeddings can be computed ahead of time. When users enter a query, only that query embedding needs to be calculated, and then compared with the embeddings already calculated for each of the documents. It’s a lot faster. So much faster, in fact, that CA is generally not used for initial retrieval, only for reranking afterwards.

However, DE models tend to be less accurate than CA models. It would be great if it was possible to transfer some of the benefits of CA models to DE models.

The Problem

It’s unclear whether the shortcomings of DEs are due to the DE model’s capacity or because of its training procedure. DEs may be overfitting.

Dual encoders can also be improved by distillation, of which there are two kinds:

1. Logit matching. Try to match embeddings between teacher and student.
2. Probability matching. Try to match the softmax probabilities between teacher and student.

Part of the cause of this discrepancy may be the fact that DE models have noisier updates. DE models may have difficulty modeling negative scores since updating their weights on positive (q, d+) pairs can inadvertently increase scores for negative (q, d-) pairs. Dropout also doesn’t seem to mitigate overfitting.

Solutions

Previous work has tried to improve training procedures in several ways:

1. Adjusting the scoring layer. Usually embeddings are scored with a simple dot product, but a more sophisticated scoring function may be able to capture more information at the cost of inference speed.
2. Distilling predictions from CA models. Model distillation uses a teacher-student framework where the smaller “student” model attempts to mirror the “teacher”. This paper explores a new approach to distillation.

The authors introduce multi-margin MSE loss (M3SE):

M3SE loss attempts to match the margins of score differences between teacher and student. For performance reasons, however, rather than matching each margin, it only encourages the student to be less than or equal to the teacher’s highest negative score.

M3SE can be seen as an extension of Margin MSE loss where instead of matching logits it matches raw scores. It can also be seen as a smooth approximation of softmax cross-entropy loss. The authors also highlight parallels between M3SE and RankDistil.

Results

The TLDR on opponent modeling is that if we have a pretty good idea of what the opponent might do, we can beat them faster by playing moves which aren’t objectively “optimal” as far as minimax is concerned. Here, Maia 1900 is a model of a relatively high-level chess player. Antimaia 1900 is specifically designed to counter Maia 1900.

• Maia 1900 uses Maia Chess weights trained on games from players rated near 1900.
• Antimaia 1900 uses Maia 1900 weights and Stockfish @10.
• Stockfish @10 is Stockfish 14.1 evaluated at depth 10.

As you can see, Antimaia finishes its games much more quickly (12-16 moves faster!). This demonstrates that Antimaia is exploiting weaknesses in Maia’s policy as measured by Stockfish, and playing high risk/high reward moves which usually pay off.

Selected Games

I had to keep reminding myself that even though these games look ridiculous, Maia plays poorly because Antimaia is finding moves where Maia is most likely to mess up.

Antimaia vs Maia

• Game 15: Smothered mate comes out of nowhere.

• Game 24: Antimaia lays a poison pawn trap then waits patiently for 2 moves before Maia falls for it, blundering a queen.

• Game 32: The shortest game. Another nice checkmate delivered by the knight.

• Game 59: Maia blunders a piece and then falls for Antimaia’s poison rook to get checkmated.

• Game 75: Promoting to a knight not because it’s the best move, but because it doesn’t matter.

• Game 88: Antimaia makes the crazy move Bh7 but Maia blunders into checkmate by taking with the knight.

Maia vs. Antimaia

• Game 2: Queen takes e3 is objectively a horrible move but Antimaia gambles correctly that Maia won’t know how to respond.

• Game 7: Antimaia baits Maia into blundering a queen and mate in 3.

• Game 51: Antimaia tricks Maia with a poison pawn.

• Game 59: Maia manages to lose in 5 moves.

• Game 69: Antimaia swindles a queen.

The rest of the games

I added the rest of the games to github: https://github.com/lukesalamone/antimaia_games

If you want to fly across the Pacific Ocean, you’ll have to board an airplane and fly around 12 hours. It’s pretty slow. A much faster route would be to go directly through the center of the earth. “Digging to China” was a common expression I heard growing up, with the implication that the opposite side of the globe is somewhere in Asia.

For the US, that’s actually not true. The opposite end of the earth for most of the US is actually somewhere west of Australia. For North America, there are actually very few places whose antipode (the point on the exact opposite side of the globe) is on dry land. If you could dig a hole through the center of the earth from the US, you’d pop out in the Pacific Ocean.

It is actually possible to “dig to China” though. Most of South America (excluding most of Brazil) has an antipode in east Asia. In fact, several of the world’s biggest cities have antipodes on dry land. I’ll get to this later.

Calculating Antipodes

Given the latitude and longitude of a point on earth, it’s pretty easy to calculate its antipode:

def antipode(lat, lon):
  lat = 0 - lat
  lon = 180 + lon if lon < 0 else lon - 180
  return lat, lon

Dry land?

Once I have the antipode of a particular latitude and longitude, I want to know if it’s on dry land or not. The process of figuring out what’s at a particular latitude/longitude point is called reverse geocoding, and there are a few nice tools for doing this for dry land. These tools will give fairly good information about the country a point is in.

But for now all I care about is whether a particular point is on land or water, not what country it’s in. I wasn’t able to find a great tool online that does this, so I opted for a more brute force method.

I used this picture from Wikipedia which represents almost all water as the same color of blue. Based on that, I can create a boolean mask to indicate oceans and lakes, and reverse project pixels from that image to geographic coordinates. Simple enough.

Pixels to Geocoordinates

Once you have the (x,y) coordinates they need to be converted to longitude and latitude. Latitude is how far north (or south) you are on the globe and ranges from -90 degrees at the North Pole to 90 degrees at the South Pole. On a flat map this is represented by a y coordinate.

Longitude represents how far east or west you are from the Prime Meridian (which runs through Greenwich, UK) and ranges from -180 to 180, both of which represent the same meridian.

How a latitude/longitude point is projected onto the globe depends on the map projection. In this case we’re using Equirectangular projection, so we just need to reverse that operation to get the x/y coordinates. The formula for doing so is:

#height is the height of the image
def lat2y(lat):
    return int((90 - lat) * (height / 180))

# width is the width of the image
def lon2x(lon):
    return int((lon + 180) * (width / 360))

# the reverse operations of the above
def y2lat(y):
    return 90 - 180 * y / height

def x2lon(x):
    return 360 * x / width - 180

Pixel Antipodes

Now we can calculate the antipodes of a pixel.

def pixel_antipode(y, x):
    lat, lon = antipode(y2lat(y), x2lon(x))
    return lat2y(lat), lon2x(lon)

Antipodal Cities

If we look at a list of the world’s largest cities, which of them have antipodes on land? In the top 20, only four cities do.

Step 0: Define some constants

We will use the three tokens ['major', 'league', 'baseball'] with model dimension 4 (d_model=4) and two attention heads (num_heads=2).

num_heads = 2
d_model = 4
d_k = d_v = d_model // num_heads

Here we also define the dimensions of the key and value here so that after after concatenation in the multihead attention algorithm, the final result will have dimension (seq_len, d_model).

Step 1: Get token embeddings

Usually we would tokenize the input string “major league baseball” and then look up each corresponding embedding in the model’s token embedding table. Here, we will just use random numbers to illustrate the math.

major = array([0.4, 0.2, 0.1, 0.3])
league = array([0.5, 0.3, 0. , 0.9])
baseball = array([0.8, 0.4, 0.2, 0.6])

x = np.stack([major, league, baseball])

Step 2: Define weight matrices

Here we define four weight matrices:

• The query, key, and value weight matrices have dimensions (n_heads, d_model, d_k). It is pretty standard for d_k == d_v, although it is technically not mathematically required to be the case.
• The output projection matrix is typically a square matrix with dimension (d_model, d_model).

wq = array([
    [
        [0.5, 0.6],
        [0.8, 0.3],
        [0.7, 0.2],
        [0.1, 0.8]
    ],
    [
        [0.3, 0.6],
        [0.5, 0.8],
        [0.9, 0.1],
        [0.5, 0.8]
    ]
])

wk = array([
    [
        [0. , 0.7],
        [0.4, 0.3],
        [0.8, 0.3],
        [0.7, 0.4]
    ],
    [
        [0.4, 0.9],
        [0.1, 0.7],
        [0.6, 0.6],
        [0.1, 0.6]
    ]
])

wv = array([
    [
        [0.4, 0.5],
        [0.7, 0.6],
        [0.9, 0.8],
        [0. , 0.3]
    ],
    [
        [0.9, 0.2],
        [0.2, 0.8],
        [0.7, 0. ],
        [0.1, 0.4]
    ]
])

wo = array([
    [0.2, 0.2, 0.6, 0.1],
    [0.3, 0.2, 0.2, 0.8],
    [0.2, 0.5, 0.1, 0.3],
    [0.6, 0.9, 0.5, 0.7]
])

Step 3: The multi-head attention algorithm

Now that we have our pieces ready, we can compute multihead attention on the input. The following is an example to illustrate the mechanism of multihead attention. Real implementations are vectorized to avoid slow loops.

# the input x has shape (seq_len, d_model)
def mha(x):
    heads = []
    for h in range(num_heads):
        q = x@wq[h]
        k = x@wk[h]
        v = x@wv[h]
        head = sdpa(q, k, v) # head has shape (seq_len, d_v)
        heads.append(head)
    concat = np.concatenate(heads, axis=1)
    return concat @ wo

As you can see, for each head we:

• compute the query matrix by matmul of input x with the query weights for that head. Notice that in this operation the input x has shape (seq_len, d_model) and the query weights for that head will have shape (d_model, d_k), so the resulting matrix will have shape (seq_len, d_k).
• compute the key and value matrices in the same way
• compute the final value of the head after scaled dot product attention between that head’s query, key, and value matrices. We will talk about SDPA in the next step
• The final value of the head will have shape (seq_len, d_v). If we think of these dimensions as rows and columns, the “head” is really operating on a group of columns to make a “skinny” matrix.

Finally, we concatenate all heads over the last dimension. This combines all “skinny” head matrices into a final matrix with shape (seq_len, d_model). We then matmul that combined matrix by the out projection matrix.

Step 4: Scaled dot product attention

Scaled dot product attention here is the same as in single-head attention (i.e. self-attention). The difference here is that we are sending “skinnier” matrices into the SPDA function.

def sdpa(q, k, v):
    raw_attn_scores = q@k.T
    norm_attn_scores = raw_attn_scores / d_k ** 0.5
    weights = np.stack([
        softmax(row)
        for row in norm_attn_scores
    ])
    return weights @ v

First, we matmul the query by the transpose of the key. Both the query and key matrices have dimension (seq_len, d_v), so raw_attn_scores will have dimension (seq_len, seq_len).

Next, we divide raw_attn_scores by the square root of d_k. This helps to stabilize training since the variance here scales with d_k, specifically in the softmax in the next step.

After that, we compute the softmax values for each row. Since each corresponds to a token in the input, this can be thought of as emphasizing which tokens in the sequence are most relevant to each query token. This operation doesn’t change the input shape.

Finally, we matmul the weights with the value matrix, (seq_len, seq_len) x (seq_len, d_v) resulting in a final matrix of (seq_len, d_v).

In this post I will describe the attention mechanism, commonly used in transformers, a popular neural language architecture. Most of the most well-known large language models of late are based on the transformer architecture. Attention was first described in Attention is All You Need by Vaswani et al.

What is attention?

At a high level, attention is a mechanism for neural networks to boost portions of an input which are relevant and ignore those which aren’t. In language models, attention is used as a way for the model to learn which portions of a sentence are relevant to each word.

What is attention for?

Let’s use an example:

I am sitting at the library with my friend.

It should be pretty clear that not all words in this sentence are equally important. What words are relevant to “I”? Probably “sitting”, “library”, and “friend”. Likewise, “the” might only be relevant to “library”. Attention provides a way for a model to increase and decrease the importance of each word.

Since the value of each token in the sequence is dependent on other tokens, this method of generating word embeddings is very different from more classical methods like Word2Vec and GloVe. There is no one fixed vector for a given word.

And this makes sense. Many words in English are homonyms, and have identical spellings for distinct meanings. For example, “rock” is a genre of music but also can mean a stone. The word “run” has 645 meanings and has recently replaced “set” as the word with the most definitions. It would not make sense for all of these homonyms to have the same vector.

An interactive example

You can hover over each word to see the relative importances of each word in the sentence to the hovered word.

How does self-attention work?

In linear algebra terms this means multiplying the 1xd input vector by a matrix of size dxd. In neural network terms, this means passing the input vector through a full-connected layer. After M_k and M_q are multiplied, they are normalized by the square root of d_k, a constant representing the dimension of the input vector.

Can we attend to multiple parts of a sentence?

Yes, that is called multi-headed attention. Its architecture is very similar, using additional M_k, M_q, and M_v matrices for each additional “attention head”. In the Vaswani paper they used 8 heads.

How do transformers compare with other architectures (e.g. RNN/CNN)?

Generally speaking, RNNs are able to memorize but not parallelize, and CNNs are able to parallelize but not memorize. Transformers are able to do both.

The Vaswani paper outlines three main benefits:

1. Computational complexity per layer. Self-attention layers are faster than recurrent layers when the input sequence length is smaller than the input vector dimensionality.

2. The opportunity to parallelize calculations. Each head in multi-headed attention can be computed separately in an encoder layer.

3. Easier to learn long-range dependencies. For many English sentences, especially fairly complex ones found in more scientific writings, the full context of a word cannot be learned from its immediate neighbors. Sometimes it can’t even be found in the same sentence. However, even though most language models prior to the transformer had theoretically infinite input sequence lengths, in practice it was quite difficult for them to learn long-range dependencies. Because a transformer sees its whole input simultaneously, Vaswani argues, it is able to more easily learn those dependencies.

Is that all?

No. In part two I will describe the encoder and decoder blocks, as well as the self-supervised training process.

1. The search needs to be fast, useful, and aesthetically pleasing.
2. Search in the browser. Standing up a server is a lot of extra work. It’s also overkill since I only have about 30 articles so far.

Semantic search

I ran some experiments with small neural networks deployed using ONNX but ultimately they didn’t seem to be a good fit for this blog. The search experience was not quite as snappy as I’d have liked it to be, and while I was able to get the model under 10MB, it still added a good amount of bloat to the page size. Further, it wasn’t clear to me that the search results were significantly better, and in some cases they were worse. In any case, the advantages were not enough to justify the added page size.

I learned a lot with this experiment which probably deserves a blog post of its own at some point. I might revisit semantic matching at a later point as well. (Post script: I did.)

BM25

Rather than loading a neural network, I decided instead to index the text of my blog and use BM25 to rank documents.

One difficulty in this was finding the actual algorithm which didn’t include a ton of extra code. I was only interested in the function itself, not all of the extra boilerplate code floating around. In the end I broke down and rewrote it myself.

Tokenizing Text

BM25 works by scoring document tokens against query tokens. Documents with more relevant tokens should be scored higher than those with fewer relevant tokens.

This raises the question of how to convert a query string into a group of tokens. Initially, I simply split by word. That is,

"hello my friend" => ['hello', 'my', 'friend']

Although this works pretty well, it has the downside that different forms of a word will be missed. For example “friends” will not match with “friend”. To solve this problem, I decided to use fixed token sizes of three characters instead:

"hello my friend" => ['hel', 'ell', 'llo', 'lo#', '#my', 'my#', '#fr', 'fri', 'rie', 'ien', 'end']

In this way “friend” will be tokenized very similarly to “friends”, so the penalty for mismatches between forms of words will be minimized. Subjectively, I felt that this yielded a much more fluid and versatile ranking as well.

Multi-factor scoring

My blog posts have title, body, and tags. I could simply concatenate them all together and score the combined string, and it would work pretty well. But this seems to throw away the fact that the title and tags may have more concentrated information than the article body does. They deserve to be weighted differently. I decided to weight the title/body/tags components as 30/50/20 in the final ranking.

In the previous post I discussed the history of chess engines and why they don’t “think” like we think. Trading interpretability for computation cycles ultimately led to the engines we have today, fairly alien in nature and perhaps less pedagogically useful because of it. At the time, though, the goal was to beat human grandmasters by any means necessary, a great engineering feat that the field had been working on for decades.

This post contains two related proposals. The first is a chess engine tournament, unique in the type of engine which will be permitted to enter and likely to succeed. Importantly, the vast majority of engines currently holding the highest performance ratings will likely not be effective.

The second proposal is the outlines of a chess engine that is likely to be successful in this tournament, taking advantage of the highly qualitative nature of chess position evaluation. Although it is unlikely to perform as strongly against top-performing engines, there are several distinct advantages of such an engine. In short, there is likely to be a great deal of educational value as well as financial incentive driving the construction of highly successful qualitative chess engines.

Qualitative chess analysis

Qualitatively, there are many aspects to a chess game that may be captured. Let’s take a look at the way a grandmaster analyzes a position. It will become quite apparent that at the highest levels, the qualitative aspects of position analysis dominate over quantitative aspects (i.e. the number and value of each piece).

In the selected lecture, Grandmaster Varuzhan Akobian details a game he played previously. At a key moment of the game, Akobian sacrificed his rook for a key pawn in the center of the board. The resulting position is reproduced for convenience in Figure 1.

His analysis starts at 25:15 in the video:

I would like you to spend a minute or two just to give me the evaluation of this position. It may not seem that clear because I’m down the exchange. [Novice chess players are taught that chess pieces have quantitative values, which may come into consideration when exchanging one piece for another. These values are measured in terms of pawns. Knights and bishops are generally understood to be worth 3 pawns, rooks are worth 5, and queens are worth 9.] I have a knight and a pawn for a rook. Rook is valued 5, knight and a pawn is 4. It may seem like I’m down a pawn here. But what do you think is the proper evaluation of this position?

…Basically white is very active. There are a few other things we can mention about white’s position, that it’s very strong. What else is very strong? White’s king is very safe, he cannot attack me. But how about the black king? Do you think the black king is very safe? [No.] For example, I could put my queen here [the e4 square] then I have a battery! Remember when we have a queen and a bishop on the same diagonal we call that a battery. And suddenly if I can deflect this queen [black queen on the g7 square] I will just go queen takes pawn, checkmate!

His dark square bishop is basically trapped behind his own pawns so it’s ineffective. . . . My bishop is very active. . . . And one more thing that you can mention. Passed pawn, exactly! And it’s a very strong passed pawn because with a knight on d6 very quickly [the pawn] will turn into [a queen].

. . . How much advantage does white have here? Big advantage, slight advantage, maybe winning? . . . We’re not going to use Houdini [a chess engine], Houdini will probably say black is slightly worse. But in practical play, I would be very comfortable to play this against anybody, and pretty comfortable I can win this position for white.

Note that quantitative analysis is almost entirely absent from GM Akobian’s evaluation. Towards the beginning he mentions that he has sacrificed his rook for a knight and a pawn, and consequently is at a material deficit. However, he quickly discards this shallow evaluation, going so far as to label his subsequent qualitative evaluation as the “proper” evaluation.

GM Akobian goes on to mention several other qualitative features of the position which are difficult to assign quantitative value to. Firstly, the activity of his pieces means that it is much easier to play the position as white because his pieces are on better squares, including some deep in black’s half of the board. The lack of activity is mentioned later on, noting that black’s bishop is essentially trapped behind his own pieces.

King safety is another difficult thing to quantify. In the given position, it is difficult to find a way that black can even check the white king. Moving the d8 rook to b1 will take 2 moves, and even then the b1 square is guarded by the bishop on d3. So the white king is indeed quite safe. In contrast black king is quite vulnerable, guarded mainly by the black queen, who is herself vulnerable to deflection or direct attack. (Deflection is a chess tactic which involves “distracting” an opponent’s piece which plays an important defensive role. For example, a piece which is defending two pieces simultaneously may be deflected by capturing one of the defended pieces.)

GM Akobian emphasizes the weakness of black’s king by sketching out a simple game-winning checkmate plan: arrange the bishop and queen in a battery which attacks the h7 pawn, deflect the black queen, and deliver checkmate with the queen by taking the h7 pawn. Although it is not immediately clear how to implement the plan, this type of simple plan creates a well-defined long-term “threat” that black must contend with.

Another threat he mentions is encompassed by white’s passed pawn on c5. This pawn may become a queen, which would become an insurmountable advantage for white. Therefore, this threat is another long-term vulnerability for black. (A passed pawn is a pawn which cannot be stopped or attacked by an opponent’s pawns. This occurs when there are no opponent pawns in the “file” (vertical column) of the pawn, as well as the file to the left and the right, if applicable. For example, a pawn in the C file is a passed pawn if there are no opponent pawns in the B, C or D files. A pawn in the H file is passed if there are no opponent pawns in the G or H files.)

Finally, note that GM Akobian does not assign a quantitative value to the board position, but rather a “very comfortable to win” assessment. Very little of this analysis involves quantities, but rather qualitative situations which must be dealt with. Consequently, it seems that qualitative reasoning is an ideal tool which a chess engine might use.

Qualitative analysis is more human

The nature of expert-level perception described by GM Akobian was studied directly in a 1973 paper by Chase and Simon. Participants at three different levels of chess ability (a master-level player, an experienced club-level player, and a novice) were asked to complete two chess-related cognitive tasks. The first was a perception task, requiring him to reproduce a chess position on an adjacent board as quickly as possible, with the model board in plain view. The second task was a memory task, requiring participants to reproduce a position from memory after viewing it for only 5 seconds.

Importantly, the perceptual study attended to chess players’ tendencies to “chunk” the board position as they reproduced positions, tending to remember groups of interrelated pieces. These pieces tended to have relationships which the authors characterized in five ways: a piece attacks another, a piece defends another, two pieces are adjacent, two pieces are the same color, two pieces are the same type.

The results of the study found that “the C, S, and null relations are low because subjects are placing pieces which usually have multiple relations. Thus, from the within-glance relations, it appears that subjects are noticing the pawn structure, clusters of pieces of the same color, and attack and defense relations over small spatial distances.” In other words, it seems likely that human players use the qualitative relationships between pieces to remember the board.

A qualitative chess engine

It is unlikely that a qualitative chess engine will be able to entirely do away with the basic structural algorithm involved in chess calculations, i.e. minimax. We would like our qualitative engine to calculate in a way most similar to humans, and thus will require some level of ply depth to the calculations. However, a qualitative engine will have a much stronger sense of the “flow” of the game, and will thus explore fewer branches. Rather than considering each position as discrete, a qualitative engine should note how each move guides the evolution of the chess board position.

It is important to note that a qualitative chess engine may not be the most computationally efficient, a factor which was the primary motivation during the period of time when top chess engines needed to be run on supercomputers and every ounce of performance needed to be squeezed out of the machine. A qualitative engine should instead favor explainability over performance whenever possible. Instead, an engine should produce a good explanation of which moves were considered and why a particular move was chosen.

More modern qualitative research can improve upon Wilkins’ knowledge-based PARADISE approach. It is important to recognize that his knowledge base is quite similar in nature to the FAC component of the retrieval model presented by Forbus et al. in 1995, but does not take advantage of the performance speedups presented there. Because of the high number of positional examples available online, there is a huge opportunity for a performant analogical retrieval system at present. The MAC/FAC retrieval system could pay huge performance dividends in retrieval if applied to this problem.

Specifically, the Lichess database referenced above contains 1,737,489 chess “puzzles” as of the time of writing. A chess puzzle is simply a chess board position in which players are encouraged to find the best move. Each puzzle relates to one or more chess “themes” (e.g. “mate in 1”, “pin”, “discovered attack”, etc.), analogous to Wilkins’ concepts outlined above. Each puzzle also includes the best move to make in the position. Some research will need to be done to derive meaning from this best move, relating it by analogy to the current position being evaluated by the engine.

Qualitative spatial calculi may also be used to construct more psychologically plausible models of chess positions than simply noting which pieces occupy which squares, seeking to emulate the models suggested by Chase and Simon. Chess pieces have intricate relationships which can be captured, and which change whenever a piece moves to another square. Importantly, however, not all relationships are affected by the movement of a single chess piece, suggesting that performance gains may be realized by recomputing only those relationships which have changed.

It is likely that low-level piece relationships may give rise to higher level relationships and tactics. For example, the concept of “capturing the defender” arises from the concept of attacking a piece A which defends B, which works when pieces A and B are attacked by pieces C and D of opposing colors. And in the case of Figure 2 below, a defender may be “deflected” to win the piece it is defending. Defensive relationships may be thought of in a chain or directed graph, with each piece defending another and the safety of a piece being considered in relation to its connection to a defensive group.

Applications of a qualitative chess engine

There are many benefits to reopening the pursuit of qualitative reasoning in chess. The first and most clear value proposition is that qualitative reasoning is likely to serve as a more plausible model for how humans think about the game. This is evidenced by the fact that as Chase and Simon found, chess players do not “see” the whole board at once, but rather in chunks of interrelated pieces. Even if the details of human mental models differ slightly from the implementation of a qualitative reasoning engine, it will be able to provide a traceable account of its decision-making process, an important step towards explainability.

As we saw in part 1, current top chess engines reason about chess in ways that are quite contrary to human intuition. Stockfish uses full-width search, considering each move in each position without prejudice and assigning numerical values to each position. As we saw from the analysis from GM Akobian, qualitative evaluations are far more meaningful to humans.

ther chess engines approach chess in an even more alien way. Specifically, it is unlikely that any engine which makes heavy use of neural evaluation functions will model human-derived organic strategies in ways which chess players will recognize. At the far end of this spectrum is the fully neural Maia chess engine, but even Lc0’s Monte-Carlo tree search precludes consideration for cognitive plausibility.

Qualitative chess engines which are able to better reproduce the types of chess reasoning used by top human chess players are also likely to serve as better pedagogical tools for those interested in studying chess. This applies at every level, from beginner to grandmaster. The skill level of such a chess engine would be quite easily tunable simply by disabling more advanced knowledge from the knowledge base.

This is a far more natural method of “handicapping” than the search depth limitations used in current chess engines. Each piece of knowledge becomes a tunable parameter to the engine. As students learn concepts, the corresponding representations in the knowledge base could be enabled, allowing for gradual learning in a far more accessible way. In fact, it is likely true that a qualitative chess engine could outperform human grandmasters (who often teach chess to others) in this respect.

Finally, it is likely that a qualitative engine would become a key component of a first line of defense against cheating in chess. Most cheating is performed by using assistance from a chess engine during online games with unsuspecting opponents. Consulting a functionally omniscient computer program can thus provide a cheater with a theoretically insurmountable advantage.

In an interview with the Perpetual Chess Podcast, Chris Callahan of the popular chess website LiChess.org stated that the majority of employees of the website work primarily to detect cheaters and yet the problem still persists. By exploiting the difference between conventional full-width engines like Stockfish and a qualitative evaluation, those working to detect cheaters will be better equipped to detect “suspicious” moves. However, qualitative chess engines are unlikely to be able to completely replace human moderation.

Nerf the Engine!

A computer chess tournament be held between chess engines can encourage the type of reasoning and gameplay which resembles human games. However, because we are not interested in the best overall chess engine, but one which can reason like a human might, the rules of the tournament will be adjusted in several key ways to discourage brute-force computational methods. We already know that calculating millions of positions can find the optimal move. But what happens when an engine is limited to e.g. 1000 positions?

Because we expect few entrants in early iterations of this special tournament, engineering an automatic enforcement mechanism for the limitations stipulated in this document are likely to be unnecessary. Engine compliance may simply be verified through manual inspection. Future iterations may include further safeguards, potentially separating the position evaluation function and directly counting the number of invocations while arbitrating the tournament to directly verify compliance.

Position limitation

Firstly, competing chess engines will be limited in the number of board positions they can evaluate during any one move. Because human grandmasters evaluate around 100 positions before making a move, the tournament arbitration system will artificially impose this limitation on all competing engines.

This cap immediately creates an issue for full-width chess engines because of chess’ high branching factor. Were an engine to evaluate each possible move, it would perform quite poorly in board positions with many possible plies and replies available, rarely reaching a depth of more than 2 or 3. As a result, any engine which naively assesses a chess board would perform quite poorly in this setup.

The practical upshot of the position limitation is that the engine will be incentivized to gather as much relevant information about a position as possible rather than optimizing for the maximum number of positions.

Position saliency

Additionally, engines will be required to implement scheduling logic which takes the time remaining into consideration. While this creates the immediate problem of how an engine should allocate its time, it creates the ancillary challenge of evaluating a position’s quiescence. Positions which are “quiet” and have few forcing moves require less evaluation than positions in which there are many non-forcing moves.

This requirement immediately motivates a qualitative chess engine to recognize the futility of falling prey to the Horizon Effect. The Horizon Effect causes engines to waste many position calculations pursuing delaying moves which amount to hopeless rabbit trails. Instead, an engine should recognize that quiescence has to do with the relationships between pieces. Humans understand this and can quite quickly see the futility of a move and terminate their search. A qualitative analysis which takes this factor into consideration will be able to save a great deal of position calculations, behaving more like a human player.

Conclusion

Given that computers have achieved and sustained superhuman capabilities in the domain of chess, the next frontier is not in building increasingly strong engines, but harnessing the present computational power to reason about the game in ways that humans do. Qualitative reasoning can provide novel and intuitive ways to reason about previously seen moves and think about the game.

Computers that play chess, otherwise known as chess engines, have existed since at least the late 1940s. Because the game was said to require the perfect combination of planning, strategy, psychology, and calculation, chess was once thought to be an activity directly correlated with intelligence, and that only a truly intelligent computer should be able to defeat humans. However, as a recent chess.com report explains, computers are now far stronger than humans:

Human chess and computer chess are different, even at the highest levels. The best humans play at an Elo rating of 2800. “Stockfish,” the most powerful chess engine, has an estimated rating of more than 3500. In a theoretical match between World Champion Magnus Carlsen vs. Stockfish, we estimate that it is most likely that Magnus Carlsen would lose every single game—no wins and no draws.

The supremacy of Machine over Man should be marked as a great triumph in artificial intelligence, and in a broad sense it is. However, I would argue that playing skill should not be the only goal, and a truly useful engine should be interpretable as well. We’ve already “solved” the problem of playing chess at an extremely high level. Yet for historical reasons, interpretability was sacrificed for the sake of speed and playing strength. However, these tradeoffs may not make as much sense today. Part 2 discusses opportunities and clear benefits for reevaluating those decisions moving forward.

Background

The general algorithm for performing search in zero-sum perfect information games like chess is known as minimax. The algorithm attempts to find the move which minimizes the opponent’s maximum (i.e. best) move. Minimax visualizes the possible future state of the game as a tree, and the value of each node in the tree as a function of the nodes following from it. If you could compute the full game tree (in simple games like tik-tac-toe this is possible), you would be able to enumerate all possible terminal nodes, and propagate the result (1 for a win, -1 for a loss, 0 for a draw) back up the tree.

That’s not possible in chess. The game tree is too big. Instead, the value of many of the nodes needs to be estimated using heuristics so that they can be propagated upwards. Therefore a fast and accurate static evaluation function is critical. (However, since the problem is recursive, the problem of creating a perfectly accurate static evaluator is as hard as evaluating that node’s entire associated game tree.)

Minimax is often augmented with “alpha-beta pruning” to reduce the number of positions which will be evaluated. This effectively cuts the computational complexity exponent in half by removing from consideration those branches which cannot affect the final result.

Other challenges in minimax evaluation exist as well. In particular, a phenomenon dubbed the “Horizon Effect” is a peculiar failure mode of minimax searches. The Horizon Effect was first described by grandmaster and computer scientist Hans Berliner in 1975. His illustration of the problem (his Figure 1.3) is reproduced in Figure 1. Algorithms that do not account for the Horizon Effect will try to “push” bad outcomes of their search beyond their search horizon, instead opting to make hopeless moves which only serve to delay the inevitable. As humans we intuitively understand this illusion.

Branching Factor

In any given position, there may be many possible moves. By some estimates, chess has an average branching factor of around 30, with other estimates putting the number around 40. It isn’t easy to find a concrete value for the branching factor of chess. One source claims (without citation) that the branching factor may be up to 40 (see pg. 117). However, a more recent statistical analysis of 2.5 million chess games put the real number closer to 31.

In any case, it is possible to dramatically reduce the branching factor by employing “selective search”, which means excluding nodes from recursive tree search altogether. Reducing the number of possible moves that will be explored in any given position promises dramatic computational speedups, allowing the tree to be searched deeper at the expense of width. Given the slow speeds (by contemporary standards) of early computers, the allure of such a technique should be clear. It is also a more psychologically plausible way of playing, since most players quickly rule out ridiculous seeming moves and focus on the most promising ones.

Early Chess Engines

An early attempt to narrow the search tree was proposed by Berliner, who in 1975 devised CAPS-II which utilized a tree search algorithm with five total “reference levels” including ALPHA and BETA. His paper also was cognizant of the work of Chase and Simon from two years prior, recognizing the need for a bottom-up board representation. Unfortunately, his board representation was largely geometrical and included little in the way of qualitative relationships between pieces. Nevertheless, the resulting program was able to solve tactical puzzles in a quite impressive manner for the time.

Another example of the use of such a narrowed search tree is PARADISE (PAttern Recognition Applied to DIrecting SEarch). This system used a knowledge base containing 200 “production rules” to match board patterns and determine the best move at any time. An example production rule is shown in Figure 2.

Because the tree search was narrower, the system was able to search to higher depths. Still, because of hardware limitations of the time, PARADISE was extremely slow, generating only 109 nodes in 20 minutes of search time. (Current chess engines evaluate millions of nodes per second.)

PARADISE executes static analysis on a position using its many production rules with the ultimate goal of creating and executing a plan. Matched production rules post concepts to the database in addition to information about the intentions of the concept and its likelihood of success.

The program is then able to use this information to form a plan, which is a set of actions for the side to move, along with corresponding plans for each defensive alternative. Because there may be many potential alternatives at each move, this plan includes many branches.

An example of such a plan is reproduced in Figure 3 above. (An interesting sidebar about this position is that based upon Stockfish analysis, this position results in inescapable checkmate for black within 15 moves. That is, even with optimal play, black will be checkmated in 15 or fewer moves. The winning sequence begins with white moving his queen to e5, exploiting the pin on the f6 pawn from the f1 rook.)

Searching only selective lines is more difficult to implement than full-width search. Further, selective searches may miss important continuations of a position, causing the computer to select an incorrect move. It was for this reason that Berliner himself, an original proponent for the application of strict logical rules in chess, decided to seek brute-force search methods instead.

Modern Chess Engines

Stockfish is currently the #1 ranked chess engine in the world. Stockfish is open-source and performs a “full width” search on the game tree. Leaf nodes are evaluated using either a classical, hand-crafted evaluation heuristic, or more recently, a neural network evaluator called NNUE. The classical evaluation function uses a set of around 30 factors weighted empirically using a dedicated testing framework called Fishtest. Altogether, the project has around 200 contributors. Optimized for speed, Stockfish can evaluate around millions of nodes per second on a typical 4-core computer.

Neural networks may also be used more directly. Leela Chess Zero, also known as Lc0 is another open-source chess engine, modeled after DeepMind’s AlphaZero chess engine. Lc0 uses Predictor + Upper Confidence Bound tree search (PUCT) to search its game tree. Leela evaluates new nodes by iteratively choosing moves from a probability distribution until it reaches an unexplored node. At that point, Leela’s neural network estimates the node’s value and propagates that value back up the tree. PUCT is very similar to Monte Carlo Tree Search, but with game “rollouts” replaced by neural network evaluations.

Another notable engine is Maia Chess. Maia is interesting for two reasons. Unlike most other engines, Maia’s objective is to model human behavior rather than to perform optimally. This is interesting because Maia is in many cases trained on suboptimal data. The other notable feature about Maia is the unusual manner in which moves are selected: instead of performing any type of tree search at all, the engine simply returns the neural network’s static evaluation of the current board position.

We can also attempt to predict the evaluation of Stockfish directly with a neural network, as Deepmind did in 2024 and as I did using one billion positions from the Gigafish dataset.

Suffice it to say, however, that despite computation no longer being a significant limitation, prior qualitative approaches still have garnered relatively little attention.

Conclusion

Engineers have historically sought to build the most powerful chess engine possible. With the singular goal of defeating human (and now computer) opponents, explainability was an expedient but understandable tradeoff to make. In the present day, however, there is no longer a question of whether humans or computers are superior chess players. It seems quite clear that the time has come to revisit some of the tradeoffs made in the past.

In part 2 I will discuss some clear benefits of a more interpretable chess engine, and some possible routes of getting there.

TLDR: Using opponent modeling we can win 2x faster than Stockfish by playing high risk, high reward moves that Stockfish will avoid.

Minimax solves for optimal opponent play, minimizing the best move an opponent could make. But what if we knew the opponent wouldn’t make the best move? What if we knew what the opponent would do ahead of time? In that case, we could beat them faster by playing moves which take advantage of this fact, even if our move isn’t objectively the best move. Don’t play the game, play the man.

For the purposes of this problem let’s assume it is a chess algorithm but this concept can be generalized to any algorithm which can use minimax. Anyone who has played a game against another person understands the idea of “downloading” your opponent: playing against a person enough that you feel like you know what they will do. Hikaru Nakamura does.

It seems like this type of player would have many possible applications, including training for specific skill levels and even trainers personalized to an individual’s playing style.

Prior work

There has been previous work in terms of opponent modeling in chess. Maia chess has been trained to model several Elo levels of opponent play which makes it ideal to be used as a backbone in my experiments.

The Policy Function

We can think of the policy function as a model of the opponent. We will therefore be optimizing an algorithm to play against this model, rather than the optimal opponent.

The specifics of the policy function are as follows. The function takes two parameters, a board position B and an opponent skill level L. The parameter L is important because the nature of the opponent’s suboptimal play is important. All levels of chess players play sub-optimally, but the types of mistakes a beginner or intermediate-level player will make are different from a grandmaster. The algorithm should exploit this fact. In practice, this algorithm should play in a much more “trappy” way than the methodical slow grind of a typical chess engine.

Formally, at move $i$ there will be $n_i$ possible moves. For each of the $n_i$ moves, we should calculate an evaluation weighted by the opponent’s probability of playing responses to that move:

$ e_i(m_i) = p_L(M_{i+1}) \cdot r(M_{i+1}) $

Where $M_{i+1}$ is the set of possible moves the opponent can make in response to move $m_i$, $p_L$ is the probability distribution of each move in $M_i$ being made, and $r$ is some real evaluation of the board position.

Let’s use an example. Suppose that at some stage in the game the algorithm must choose between 2 possible moves. For simplicity, the opponent will have 4 possible replies after each of the two possible moves. After the algorithm makes the the first possible move, the probabilities look like this:

Opponent move distribution at position 1

$ \begin{bmatrix} 0.6 \\ 0.3 \\ 0.1 \\ 0.1 \end{bmatrix} \cdot \begin{bmatrix} 5 \\ 1 \\ 0.5 \\ -1 \end{bmatrix} = 3.25 $

Opponent move distribution at position 2

And after the second possible move, the opponent move distribution looks like this:

$ \begin{bmatrix} 0.9 \\ 0.05 \\ 0.03 \\ 0.02 \end{bmatrix} \cdot \begin{bmatrix} 4 \\ 0.5 \\ 0 \\ -2 \end{bmatrix} = 3.585 $

In these tables positions are evaluated from the algorithm’s point of view, so higher evaluation=better. Assume the evaluation function is a standard evaluation from a chess engine e.g. Stockfish.

In this example it is clear that the algorithm should choose move 2, because it is highly likely that after making this move the algorithm will have a decisive advantage.

Checkmates

All of this works quite well in the middle-game, where position evaluations are all within a reasonable range. But what happens when one of the opponent’s moves leads to checkmate? How much value should we give to the possibility that an opponent might blunder into checkmate? The above calculation only works for finite numbers. The evaluation of a checkmate is +/- infinity. Instead, let’s choose a large, finite number for checkmate.

The behavior of the algorithm depends heavily on our choice for this checkmate evaluation, the CHECKMATE_WEIGHT. If the number is too high, the algorithm will favor positions where the opponent might get checkmated, playing risky moves. If the number is too low, the algorithm will choose other branches which win less quickly. After a bit of trial and error, I settled on +10 for checkmate. This is enough for the algorithm to seek out these types of positions.

A further question is what to do once checkmate is inevitable. That is, what should the algorithm do once it will definitely win? One option is to continue to sample from the distribution, continuing to play the move which maximizes winning chances. At the beginning, this was my intuition. Unfortunately, this leads to many draws by repetition.

Instead, once the game is in a forced-win situation (“mate in N moves”) the algorithm will stop sampling from the distribution and simply play the moves which lead to checkmate. While this will inevitably lead to a win, it changes the character of the algorithm somewhat. Therefore, I have limited the lookahead to checkmates in the immediate next few moves (specifically, mate in 5 or fewer moves). Since the original goal was to create an algorithm which is used for training, this seems reasonable.

Games

I simulated 400 matches testing this opponent modeling approach, which I wrote about here.

Key findings:

• Both Antimaia and Stockfish@10 beat Maia 100% of the time
• Antimaia finishes its games much more quickly (12-16 moves faster!)
• Antimaia is exploiting weaknesses in Maia’s policy by playing high risk/high reward moves

Batch Norm

Normalizes the mean and variance for all instances of a feature in a batch. Commonly used in older architectures like CNNs by smoothing out the loss landscape and allowing for higher learning rates and faster convergence. Example implementation:

def batch_norm(features: np.array, i: int, eps=1e-5) -> np.array:
    x = features[:,i]
    mean = x.mean()
    variance = x.var()
    x = (x - mean) / np.sqrt(variance + eps)
    return x

In practice, we include a gamma and beta parameter which are learnable parameters of length F, where F is the number of features. These parameters that allow the network to undo the normalization if needed.

def batch_norm2(
        features: np.array,
        gamma: np.array,
        beta: np.array,
        eps=1e-5
    ):
        mean = features.mean(axis=0)
        variance = features.var(axis=0)
        normalized = (features - mean) / np.sqrt(variance + eps)
        return gamma * normalized + beta

See also: Layer norm

Byte Pair Encoding (BPE)

A tokenization algorithm which iteratively merges the most frequent pair of tokens. For example, let’s take the sentence snug_as_a_bug_in_a_rug as our corpus (underscores used to make things easier to see).

First, we break up the text into its individual characters:

corpus: ['s', 'n', 'u', 'g', '_', 'a', 's', '_', 'a', '_', 'b', 'u', 'g', '_', 'i', 'n', '_', 'a', '_', 'r', 'u', 'g']
vocab: ['s', 'n', 'u', 'g', '_', 'a', 'b', 'i', 'r']

After splitting by chars, we have 9 unique tokens in our vocabulary. Next, we merge the most common pair in our vocabulary, ('u','g') => 'ug' and retokenize:

corpus: ['s', 'n', 'ug', '_', 'a', 's', '_', 'a', '_', 'b', 'ug', '_', 'i', 'n', '_', 'a', '_', 'r', 'ug']
vocab: ['s', 'n', 'u', 'g', '_', 'a', 'b', 'i', 'r', 'ug']

Then, we find that the most common pair is ('_','a') which gives us the next merged token "_a":

corpus: ['s', 'n', 'ug', '_a', 's', '_a', '_', 'b', 'ug', '_', 'i', 'n', '_a', '_', 'r', 'ug']
vocab: ['s', 'n', 'u', 'g', '_', 'a', 'b', 'i', 'r', 'ug', '_a']

We can repeat this process until the vocabulary is the size we need.

See also: Wordpiece tokenization

Direct Preference Optimization

Direct preference optimization (DPO) is a contrastive learning method, often used for fine-tuning an LLM to align with human preferences.

Like other contrastive learning methods (for example those used to fine-tune encoders) DPO works with (anchor, positive, negative) triplets. In this case, the anchor is the prompt, the positive sample is the preferred completion, and the negative sample is the bad completion. By computing the probabilities the models assign to each of the responses, we can train the model to assign higher probabilities to the preferred response.

DPO is simpler, more stable, and more lightweight than more traditional methods like PPO and RLHF. It does not require a second model at training time like PPO and RLHF do. It doesn’t require a separately trained reward model at all. The model is updated by simple gradient descent rather than generating during training.

However, since DPO is technically off-policy from an RL perspective, since the samples it is scoring are not actually generated from the model.

DPO pairs technically do not need to be created based on human preferences. Any reward model will do. For example:

• If we’re training a chess model, we can use a chess engine to score potential moves
• If we’re training a code generation model, we can use the compiler’s feedback to create positive and negative samples

Distillation

Distillation is the task of training a smaller “student” model to behave like a larger “teacher” model.

Classically, this has meant directly using the teacher’s output distribution, and minimizing KL(p_teacher || p_student).

Less formally, the larger teacher model can be used as a pseudo-labeler of real inputs, and a smaller student can learn from those labels.

Entropy

When training LLMs, we want some level of the model’s uncertainty when generating tokens. To do this, we multiply the probability by the negative log of the probability.

import math
def calc_per_token_entropy(probabilities):
    entropy = 0
    for probability in probabilities:
        entropy -= probability * math.log(probability)
    return entropy

The negative log of the probability can be thought of as the “surprise” from that token: high probability tokens have very little surprise to them, and low probability tokens have surprise approaching infinity. Then, multiplying by the probability itself helps to balance the total by preventing it from exploding.

There is a connection to perplexity, which is $ e^{H(x)} $, where $H(x)$ is the entropy of the distribution.

Group-Relative Preference Optimization (GRPO)

GRPO is an RL algorithm for LLM post-training which attempts to determine how good a rollout is in order to apply a reward to it.

GRPO has some advantages over more traditional PPO. In PPO, we use a second model (the “critic”) to score rollouts, and in practice this model may on the order of the same size as the model we’re training. This is both slow and resource intensive.

In GRPO, we get rid of the critic entirely.

prompt = 'write a python function to determine if a number is prime'
group_size = 8 # usually between 4 and 16
responses = [
    generate_response(model, prompt, temp=1.1)
    for _ in range(group_size)
]

rewards = [
    calculate_reward(response) for response in responses
]

mean_reward = np.mean(rewards)
std_reward = np.std(rewards)
advantages = [(reward-mean_reward)/std_reward for reward in rewards]

# policy gradient update
log_probs = [model.compute_log_prob(r) for r in responses]
pg_loss = 0
for log_prob, advantage in zip(log_probs, advantages):
    pg_loss += -log_prob * advantage

# KL penalty prevents model from changing too quickly
ref_log_probs = [ref_model.compute_log_prob(r) for r in responses]
kl = sum(r - p for p, r in zip(log_probs, ref_log_probs)) / group_size

loss = pg_loss / group_size + beta * kl

KL-Divergence

A measurement of the “distance” between two distributions.

from math import log

def kl_divergence(p:list[float], q:list[float]) -> float:
  total = 0
  for a,b in zip(p, q):
    total += a * log(a/b) if p != 0 else 0
  return total

Note that KL-divergence is not symmetrical.

Where it shows up:

• In RLHF, it is used to measure the change between old policy and the updated policy, and prevent it from drifting too far using a KL penalty.
• In knowledge distillation, we directly train on minimizing KL(p_teacher||p_student) on the distribution of logits.

Layer Norm

This could probably be called “token norm” and avoid a lot of confusion since in practice this is often used in LLMs.

# hypothetical input embeddings for "ai is cool"
sequence = np.array([
    [0.3, 0.1, 0.4, 0.1], # "ai"
    [0.5, 0.9, 0.2, 0.6], # "is"
    [0.5, 0.3, 0.5, 0.8]  # "fun"
])

def layer_norm(sequence, eps=1e-5):
    """
    Example for learning, not vectorized.

    Here, `sequence` has dimensions (seq_length, emb_dim)
    and we are ignoring the batch dimension.
    """
    means = []
    variances = []
    for emb in sequence:
        means.append(emb.mean())
        variances.append(emb.var())
    result = []
    for emb,mean,variance in zip(sequence, means, variances):
        emb -= mean
        emb /= np.sqrt(variance + eps)
        result.append(emb)
    return np.array(result)

Today, many more modern LLMs use RMSNorm instead, which skips subtracting the mean.

See also: Batch norm

Offline RL

Models are trained from a fixed dataset of preferences rather than the rollouts they would naturally generate. This is much faster, but can be less useful if the training data distribution shifts far from the the model’s actual predictions.

Examples of offline RL algorithms include DPO, IPO, and KTO.

Online RL

Models generate rollouts, get scored, and update from that signal. Considered stronger than offline RL because the model is getting direct feedback on its own predictions. However, the scoring itself can be expensive, and gradient updates must wait for the scores to come back.

Examples of online RL algorithms include PPO, GRPO, and REINFORCE.

Parallelism

Data parallel

Replicate the same model over all devices.

1. For each batch, split the data between devices.
2. Each copy computes a forward pass, a loss, and a backwards pass to compute local gradients.
3. Each device reports its gradients.
4. Each copy is updated with the same average gradient.

Model parallel

This is used when a model is too large to fit on a single device.

The model can be split between multiple devices, with each device computing its part, handing it off to the next device, then waiting for the next batch.

This is very slow. Pipeline parallelism is an improvement to Model Parallelism.

Pipeline parallel

The idea here is to reduce wait time between devices.

Rather than a single device computing and then waiting as in model parallelism, a device finishes its forward pass on a “micro-batch” and then immediately starts working on the next “micro-batch”.

Similarly, backwards passes flow in reverse.

Tensor parallel

This is used when even a single layer cannot fit on a device.

The individual tensors (weights and activations) can be split across multiple devices to enable a huge matrix multiplications.

Reinforcement Learning from Human Feedback (RLHF)

RLHF has three steps:

1. Train a reward model on human preference pairs. This is very similar to a standard pairwise Learn-to-Rank model where we train on (input, positive, negative) triplets. You can simply add a linear layer to the top of the LLM to create R(x,y) which generates the ranking scores.

2a. Calculate the expected reward $ E[R(x,y)] $

2b. Penalize large KL-divergence with the old policy

Rejection Sampling/Rejection Fine-Tuning

The idea of rejection sampling is to generate multiple responses or “roll-outs” to an input. Typically we will increase sampling temperature Responses which do not meet some criterion are “rejected” or thrown away. This leaves the responses which do pass the criterion. If we fine-tune on these examples, this is called rejection fine-tuning.

RFT has some similarity with RL techniques like GRPO, which also involve using multiple roll-outs.

On the one hand, RFT is very simple and intuitive. It avoids a lot of the RL complexity such as entropy collapse and policy divergence. Additionally, because GRPO relies on the relative advantages within the group, RFT can outperform GRPO in situations where all members of the group have perfect or close to perfect rewards.

• If all members have perfect scores, the advantage will be zero for each completion.
• If all members have close to perfect scores, GRPO will amplify noise within the group.

On the other hand, RFT can be far less sample efficient than RL, because it only learns from positive examples (since the negatives are rejected). For hard problems, your acceptance rate for the problem may be very low.

• If your acceptance rate is 1 in 16, you will throw away 94% of your work to get the 6% to learn from.
• Meanwhile with GRPO, if 1 out of 16 samples in the group is correct, your algorithm can learn from all 16.
• However, if none of the samples in the group get a reward, the relative advantage for all members collapses to zero and you get nothing. In this case, it would be better to keep generating until we get a correct completion.

For this reason, we can use RFT until the acceptance rate is high enough for GRPO.

Speculative Decoding

The idea of speculative decoding is to use a small “draft” model to speed up inference for the big “target” model. This is possible because the target model can predict probabilities for multiple sequences in parallel. For example, a small draft model might predict:

draft_tokens = ['take', 'me', 'out', 'to', 'the', 'ballgame']

Now, the target model can simply predict the probabilities for each of these in parallel:

probabilities = model([
  {'token': 'take'},
  {'token': 'me', 'prefix':['take']},
  {'token': 'out', 'prefix':['take', 'me']},
  {'token': 'to', 'prefix':['take', 'me', 'out']},
  {'token': 'the', 'prefix':['take', 'me', 'out', 'to']},
  {'token': 'ballgame', 'prefix':['take', 'me', 'out', 'to', 'the']},
])

Because we’re not just using greedy decoding, we need to accept tokens from the draft model probabilistically. In other words, we accept a token with probability p_accept(token) = min(1, target_prob(token) / draft_prob(token)).

• So if the target model thinks the token is a lot less likely than the draft model does, the ratio will be small and p_accept will be small for that token.
• On the other hand, if the target model agrees with the draft model the ratio will be close to 1 and p_accept will be close to 1.
• In the case where the target model thinks the token is more likely than the draft model did, it’s somewhat of a happy accident that the token was actually selected by the draft model but the acceptance probability will be 1.

If the target model rejects a token that the draft model proposes, none of its subsequent tokens will be considered. The draft model will rerun with all of the accepted tokens and propose new tokens.

This is how it is possible to speed up the target model’s inference without losing quality.

Wordpiece Tokenization

An algorithm for training a tokenizer with a slightly more advanced rule for merging tokens. For example, consider the following pair frequencies in a corpus abracadabra:

Counter({
    ('a', 'b'): 4,
    ('b', 'a'): 3,
    ('a', 'c'): 2,
    ('d', 'a'): 2,
    ('c', 'a'): 1,
    ('c', 'd'): 1,
    ('a', 'd'): 1,
})

Under byte-pair encoding tokenization, we would automatically merge ('a','b') => 'ab'. However, wordpiece tokenization also takes the frequencies of the tokens into account individually as well:

Counter({
    'a':7,
    'b':4,
    'c':2,
    'd':2
})

In wordpiece tokenization, we merge the tokens when the pair maximizes score(x,y) = freq(xy) / (freq(x) * freq(y)). In other words, we penalize pairs if they are made up of common tokens. Let’s compute the scores:

scores = {
    ('a', 'b'): 4÷(7*4) = 0.14,
    ('b', 'a'): 3÷(4*7) = 0.11,
    ('a', 'c'): 2÷(7*2) = 0.14,
    ('d', 'a'): 2÷(2*7) = 0.14,
    ('c', 'a'): 1÷(2*7) = 0.04,
    ('c', 'd'): 1÷(2*2) = 0.25,
    ('a', 'd'): 1÷(7*2) = 0.04,
}

In this case, ('c','d') is actually the winning pair because it is composed of rare tokens.

Chess has been played in different forms since the seventh century, and in its modern form since the nineteenth century. Opening theory, i.e. the study of the best moves to begin the game with, has been developing since then.

The EGG Opening

I was inspired the other day after watching a video by Eric Rosen called “Quadruple Egg”. Chess boards are usually notated from left to right with the letters A through H. (These columns are called “files”.) The “egg” opening, therefore, involves moving the E pawn, then the G pawn twice, spelling the word “egg”. It’s an extremely unconventional opening, but how bad is it compared with unrestricted openings?

What are the best (i.e. least bad) alphabetic openings?

We can call the Egg Opening one version of an alphabetical chess opening: an opening where the piece moved must start on a file corresponding to the next letter in a given word. In this definition, the pieces don’t have to be pawns.

This raises an interesting question: what are the best and worst alphabetical openings? Under engine evaluation, they will all be losing, because the opponent is not limited in which piece they can move. However, as we will see, some are much worse than others.

Four letters

English doesn’t have an official dictionary, so I chose this list of four letter words. From there, we need to eliminate all words with letters after H. This leaves us with 37 words. Don’t ask me what they all mean:

abbe, abed, aced, ache, aged, agha,
baba, babe, bach, bade, bead, beef,
cafe, cage, ceca, cede, chad, chef,
dace, dada, dead, deaf, deed,
each, edda, edge, egad,
face, fade, feed,
gaff, gaga, gage, geed, ghee,
head, heed

A simple way of measuring how bad each of these is would be to evaluate them against Stockfish to see how badly we’re losing after playing these moves. In this simulation, we play as white each time.

This measurement approach has a problem, though: Stockfish is assuming that white can make any move, and playing black’s moves accordingly. In reality, white’s moves are extremely restricted. If black knew about this restriction, he would be much more aggressive.

Further, if black were to discover which word white is playing, he could punish white. And even if black doesn’t know exactly which word white is playing, black could eliminate moves which definitely can’t form English words. For example, if white’s first move is H, the next move will definitely be E.

Alphabetical Chess

What if both players were required to play their first four moves from dictionary words? And what if those words were pre-assigned but secret to each other? For example, if each player drew a card from a deck. In this game, if a player is unable to make a move with a piece of the correct letter, they lose.

There are a few interesting features of this chess variant:

1. Your starting word will have a big impact on how the game will turn out.
2. Figuring out your opponent’s word will give you a big advantage.
3. Checks and other attacks can be devastating. You may not be able to respond to the attack, and checks may end the game immediately.

In that case, it isn’t necessarily true that the best and worst words from above will still be the best and worst when pitted against another randomly chosen word. It’s a rock-paper-scissors situation: depending on which word A you happen to get, there is another word B which best counters your word, and another word C which your word best counters.

Here the y-axis shows the word which was played with the white pieces, and the x-axis shows the word played with black.

Cool Things About Alphabet Chess

There are a few interesting benefits to this variant:

1. Handicaps are quite easily built in. The longer your secret word is, the more restricted your play is and the more of a handicap you will have.

2. An element of gamesmanship and imperfect information. Players can win on the board with good moves, but also “off-board” by figuring out their opponent’s secret word.

3. Prevents good players from using memorized openings.

Play Alphabet Chess

If you’d like to try it out with a friend and a chessboard, here’s an online tool that’ll pick a random word: https://alphachess.surge.sh

Paper link: https://arxiv.org/abs/1906.05317

This paper describes COMET, a method of generating knowledge bases automatically. Previous work largely focused on encyclopedic knowledge, which has well-defined relationships. This paper, however, focuses on commonsense knowledge. In this paper the authors introduce a “commonsense transformer” which trains on a knowledge base consisting of tuples and a pre-trained language model. Their trained model generates new nodes in the knowledge graph and completes phrases based on edges in the existing graph.

Training

The training data consists of tuples in the form of {subject, relation, object}, and the task is to generate the object given the subject and relation. The model architecture was the Vaswani GPT. The architecture is fairly straightforward for a transformer. The only difference of note is the representation of the knowledge tuple as a concatenated sequence of the words of each item of the tuple, along with a position embedding.

To train, the researchers train COMET with Atomic and ConceptNet as knowledge seed sets. The Atomic dataset contains 877k tuples of social commonsense knowledge. Experiments with Atomic used bleu-2 as an automatic evaluation metric. The researchers also used Amazon Mechanical Turk to identify whether the generated responses were reasonable.

The researchers also tested on ConceptNet. Relations were in the same {subject, relation, object} format. This dataset was evaluated using perplexity on the test set and accuracy of the generated positive samples on the pre-trained Bilinear AVG model.

Results

Atomic experiments showed that the model is able to generate coherent knowledge even when only 10% of the data in a domain is available. The researchers found that 91.7% of the greedily decoded ConceptNet examples were correct. The model seems to be producing fairly good knowledge tuples. However, sometimes the generated tuples are merely simplifications of previous knowledge.

from torch.utils.data import Dataset, DataLoader

class CustomDataset(Dataset):
  def __init__(self, features, labels):

    assert len(features) == len(labels)
    self.features = features
    self.labels = labels

  def __len__(self):
    return len(self.features)

  def __getitem__(self, idx):
    return self.features[idx], self.labels[idx]

Next, create a custom dataloader where we specify the batch size.

features, labels = load_data()

# features & labels must have equal lengths
# e.g. features = [[1,2,3],[4,5,6]]
#      labels = [7,8]

dataset = CustomDataset(features, labels)
dataloader = DataLoader(dataset,
                        batch_size=batch_size,
                        shuffle=True)

Finally, iterate over the dataloader during training.

for epoch in range(num_epochs):
  for x, y in train_dataloader:
    # do stuff

tar -xf archive.tar.gz

If you want to compress a directory

tar -czvf archive.tar.gz /path/to/directory

That’s all.

Grover functions in an adversarial manner:

• an “adversary” generates synthetic stories
• a “discriminator” identifies fake stories

Current generative models are fairly good at creating realistic-looking text, but largely lack the ability to be controlled via tunable parameters. In contrast, Grover models news stories as distributions like (domain, date, authors, headline, body) to be sampled from. Adversaries specify domain, date, and headline, and the generator creates the body and author as well as a more appropriate headline.

Grover is built using a transformer architecture similar to BERT and GPT. Depending on the model size, the number of parameters varies from similar to GPT to on par with GPT-2. The authors use the RealNews corpus, resulting in 120 gigabytes of total file size. Training took 2 weeks on 256 TPU v3 cores.

The results of the authors’ subjective experiments show that humans have a hard time identifying Grover-written propaganda. Using Mechanical Turk, articles were rated according to stylistic consistency, content sensibility, and overall trustworthiness. Grover, on the other hand, turns out to be a fairly good discriminator.

The authors also tested GPT2, BERT, and FastText on the task of classifying news as human or synthetic. The researchers set up the experiment in an unpaired setting (give a human/synthetic classification for a single article) and a paired setting (determine which is the human and which is the synthetic between 2 articles). Unsurprisingly, the paired setting was far easier. Also unsurprisingly, the larger models perform better at discriminating when paired with smaller generators.

When comparing generator/discriminator pairs with the same numbers of parameters, classification accuracy was around 90-92%. Grover tends to leave artifacts when generating text, and this fact may be part of the reason discriminators are so good at identifying synthetic text. For one, the authors identify “exposure bias” as one of the reasons. The fact that Grover is never trained on generated text, only on human-authored articles seems to contribute to this. Perplexity also tends to vary over the length of the generated article, and depending on the sampling variance may fall out of the distribution of human language.

Create an ssh tunnel to your remote machine:

ssh -L 8080:localhost:8080 user@12.34.56.78

# or use a .pem file to connect to ec2
ssh -L 8080:localhost:8080 -i "aws.pem" ec2-user@ec2-12-34-56-78.compute-1.amazonaws.com

Start Jupyter on that machine in headless mode:

jupyter notebook --no-browser --port=8080

Use a browser to open one of the urls that Jupyter presents:
http://localhost:8080/?token=xyz

(In this table we’re keeping track of the number of days. If you want probabilities, divide each value in the table by 180. But I think whole numbers are easier to think about so I’m keeping them.)

To marginalize one of the variables, we just sum one of the variables. For example, to marginalize the weather, we would sum each of the rows to find that 96/180=53% of the time tennis was played, and 47% of the time tennis was not played. Likewise, to marginalize the boolean variable of whether tennis was played, we just sum the columns: no matter whether tennis was played on that day, how many days was it sunny? 140.

Another way of saying this is the marginal distribution of sunny weather is the first column (containing 70 and 70). The marginal distribution of playing tennis is the first row (containing 70, 25, and 1).

from google.colab import drive
drive.mount('/content/gdrive')

Follow the prompts from there. That is all.

The code used in this experiment can be found on my Github.

BERT

The Devlin et al. model was released in November 2018. It is a transformer-based language model pretrained on masked input (also known as the cloze task). During pretraining, 15% of tokens are hidden from the model, and it is trained to predict the masked tokens. As a result, I was able to evaluate its ability to correctly predict a masked token at a random position in a fixed-size input.

I looked at the following varieties of BERT:

This table also includes corresponding GPT-2 models which have a similar number of parameters. Source

GPT-2

The Radford et al. model hit the scene in February of 2019. Like BERT it is a transformer-based model, and comes in various sizes ranging from 117M parameters up to 1.5B parameters (gpt2-xl). Because GPT-2 is an autoregressive model, experiments with this family of models perform one token of generation following input context, comparing with the target token for accuracy measurement.

Here we will be evaluating two flavors of this model:

This table also includes corresponding BERT models which have a similar number of parameters. Source

Wikitext Token prediction

To evaluate the models, I sampled 10,000 random sequences from Wikitext-2.

For BERT, a random sequence of 100 tokens is selected. Then, for each sequence, a random position within that sequence is selected and masked. BERT will be required to predict this token, so accuracy is measured as the percentage of the time which its masked token is predicted correctly.

For GPT-2, a random sequence of 100 tokens is selected. Then, for each sequence, a random position within that sequence is selected. Because GPT-2 is autoregressive, it cannot attend to tokens on the right, so the sequence is truncated at the selected position. The sequence is then padded appropriately to maintain a fixed sequence length of 100.

Below we can see the performance of all 6 models on these tasks. The data has been smoothed by bucketing into groups of 5 positions at once (i.e. positions 0-4, 5-9, etc). You can see that performance of GPT-2 continues to rise as it is given additional context, while BERT models are relatively stable after being given around 5 tokens of context. Interestingly, BERT performance drops off quite steeply over the last 5-10 token positions.

When we zoom in on the final 10 positions, things start to get interesting. Both varieties of GPT-2 actually beat out all varieties of BERT at the final position.

Conclusion

BERT and GPT-2 perform quite differently on the token prediction task depending on the position of the token being predicted. For a fixed sequence length of 100 tokens, BERT performs best when the masked token is between positions 5 and 95, while GPT-2 tends to continually improve as context length increases. Interestingly, when the final token in the sequence is to be predicted, BERT’s performance falls off dramatically, while GPT-2 performance remains stable.

What is tokenization?

It’s important to understand that GPT-2 doesn’t work with strings directly. Instead, it needs to tokenize the input string, which is essentially a process for converting the string into a list of numbers, or “tokens”. It is these tokens which are passed into the model during training or for inference. As a concrete example, let’s look at a few sample sentences:

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
tokens1 = tokenizer('I love my dog')

When we look at tokens1 we see there are 4 tokens:

{'input_ids': [40, 1842, 616, 3290], 'attention_mask': [1, 1, 1, 1]}

Here what we care about is the 'input_ids' list. We can ignore the 'attention_mask' for now. We can convert the tokens in [40, 1842, 616, 3290] back into strings using tokenizer.decode:

tokens1 = tokens1['input_ids']

[tokenizer.decode(x) for x in tokens1]
# prints ['I', ' love', ' my', ' dog']

[tokenizer.decode(x).strip().lower() for x in tokens1]
# prints ['i', 'love', 'my', 'dog']

This process allows us to recover the tokens as strings from the tokenizer. For dictionary lookups, we’ll also lowercase the strings and remove the whitespace from them.

Now, let’s see what happens when we do the same thing with more complex words:

tokens2 = tokenizer('My favorite color is chartreuse')['token_ids']
[tokenizer.decode(x).strip().lower() for x in tokens2]
# prints ['my', 'favorite', 'color', 'is', 'chart', 're', 'use']

Because “chartreuse” isn’t in GPT-2’s vocabulary, it is tokenized as “chart”, “re” and “use”.

About that attention mask

For brevity I glossed over what attention_mask does above. If you’re interested in attention masks, I have a blog post on that very topic!

English words

Now it would be interesting to see how many tokens in GPT-2’s vocabulary are actually English words. This is an imprecise metric since it depends heavily on which dictionary we use. (There is no single authoritative source of all English words.) I’ll use several dictionaries and compare the results.

Enchant

PyEnchant contains a python module enchant which we can use to check if a word is spelled correctly. It can also make spelling suggestions for incorrectly spelled words:

import enchant
d = enchant.request_dict("en_US")
d.check('Hello')
# prints True

d.check('Helo')
# prints False

NLTK words

The popular NLP library NLTK also contains a word list, accessible through its corpus module.

from nltk.corpus import words

nltk_words = set(words.words())
len(nltk_words)
# prints 235892

English 350k

This list of words was taken from this github repository. It is a convenient list of lowercased words containing only letters. It seems to be the biggest of the word lists.

Lemmatization

We can bump our numbers up slightly through lemmatization:

In many languages, words appear in several inflected forms. For example, in English, the verb ’to walk’ may appear as ‘walk’, ‘walked’, ‘walks’ or ‘walking’. The base form, ‘walk’, that one might look up in a dictionary, is called the lemma for the word. The association of the base form with a part of speech is often called a lexeme of the word.

For our lemmatizer we will use WordNetLemmatizer from nltk.stem.wordnet.

Testing GPT-2 tokens

So of the tokens which GPT-2 uses, how many are English words? We can break this down metric by the dictionary used.

† indicates words were lemmatized

So the English370k word list seems to capture the most tokens from the three dictionaries. Also note the mild impact of lemmatization: although it may bump some of the percentages up a bit, it’s not enough for one dictionary to outperform another.

Looking at the tokens which aren’t in the dictionary, around 73% of them are non-word alphabetical strings. The final 27% is accounted for by symbols, numbers, and non-ascii character sequences (unicode characters from languages like Arabic, Korean, and Chinese). If we remove these, we end up with about 10k tokens containing only letters, which is around 21% of GPT-2’s total vocabulary. I’ve included this list in a github gist (duplicates removed).

Now what?

Looking at these non-word alphabetical strings, it’s interesting to see how the Internet (as GPT-2 saw it) was encoded. Then again, it also contains a lot of proper nouns which wouldn’t be in a normal dictionary like “starbucks”.

Other tokens are clearly vestiges of the scraping process which was used to gather text which GPT-2 trained on. Tokens like “rawdownloadcloneembedreportprint”, “buyableinstoreandonline”, “randomredditorwithno”, and “itemthumbnailimage” contain next to zero semantic value and the vocabulary space would probably have been better served with more meaningful tokens.

The following are the longest non-dictionary tokens found in GPT-2’s vocabulary:

We may also be able to measure performance of GPT-2 on certain tasks based on how many of the tokens were dictionary words. It might be true, for example, that sentences with higher proportions of dictionary word tokens would perform better on sentence completion tasks.

If you want to perform inference with transformers one sequence at a time, you can ignore attention masks. The “slow way” will be sufficient for your needs.

The slow way

We can tokenize and perform inference with GPT-2 one sequence at a time, but it’s slow:

from transformers import GPT2LMHeadModel, GPT2Tokenizer

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
gpt2 = GPT2LMHeadModel.from_pretrained('gpt2')

context = tokenizer('It will rain in the', return_tensors='pt')

prediction = gpt2.generate(**context, max_length=10)
tokenizer.decode(prediction[0])
# prints 'It will rain in the morning, and the rain'

It’s way faster to batch the inputs, which means adding their token vectors to the context and performing inference only once.

The un-slow way

The cool way to perform inference on many samples is with batching. It’s much faster but it’s also slightly more complicated.

tokenizer.padding_side = "left"
tokenizer.pad_token = tokenizer.eos_token

sentences = ["It will rain in the",
            "I want to eat a big bowl of",
            "My dog is"]
inputs = tokenizer(sentences, return_tensors="pt", padding=True)

output_sequences = gpt2.generate(**inputs)

for seq in output_sequences:
    print(tokenizer.decode(seq))

What’s happening here? And what does this have to do with attention masks? First let’s explain padding, then take a look at the code line by line.

We feed tokens into transformer-based language models like GPT-2 and BERT for inference as tensors. A tensor is like a python list but with a few extra features and restrictions. Specifically, for a tensor of dimension 2+, all vectors in that dimension need to be the same length. For example,

from torch import tensor

tensor([[1,2], [3,4]])  # ok
tensor([[1,2], [3]])   # error!

When we tokenize an input, it it will be turned into a tensor containing sequence of integers, each corresponding to an item in the transformer’s vocabulary. Here is an example tokenization in GPT-2:

Suppose we wanted to include a second sequence in our input:

Because these two sequences have different lengths, we can’t just combine them in one tensor. Instead, we have to pad the shorter sequences with dummy tokens so that each sequence is the same length. And because we want the model to continue to add to the right side of our sequence, we will pad the left side of shorter sequences.

This is where the attention mask comes in. The attention mask simply shows the transformer which tokens are padding, placing 0s in the positions of padding tokens and 1s in the positions of actual tokens. Now that we understand that, let’s look at the code line by line.

tokenizer.padding_side = "left"

This line tells the tokenizer to begin padding from the left (default is right) because the logits of the rightmost token will be used to predict future tokens.

tokenizer.pad_token = tokenizer.eos_token

This line specifies which token we will use for padding. It doesn’t matter which one you choose, but here we’re choosing the “end of sequence” token.

sentences = ["It will rain in the",
            "I want to eat a big bowl of",
            "My dog is"]

These three sequences all have different lengths when tokenized, so should be a good test of our padding method.

inputs = tokenizer(sentences, return_tensors="pt", padding=True)

Now we tokenize. We’re passing in the sentences from above, telling the tokenizer to use PyTorch tensors (rather than Tensorflow), and telling the tokenizer to add padding for us. We can print inputs here to confirm that, yes, tokenization is working as we thought:

{'input_ids': tensor([
    [50256, 50256, 50256,  1026,   481,  6290,   287,   262],
    [   40,   765,   284,  4483,   257,  1263,  9396,   286],
    [50256, 50256, 50256, 50256, 50256,  3666,  3290,   318]
  ]),
'attention_mask': tensor([
    [0, 0, 0, 1, 1, 1, 1, 1],
    [1, 1, 1, 1, 1, 1, 1, 1],
    [0, 0, 0, 0, 0, 1, 1, 1]
  ])}

As you can see, the first and third sequence include padding at the beginning, and the attention_mask parameter marks the position of this padding.

Now let’s actually pass this input into the model to generate new text:

output_sequences = gpt2.generate(**inputs)

If you’re unfamiliar with **kwargs syntax for function calls, this passes in the inputs dict as named parameters, using the keys as the parameter names and the values as the corresponding argument values. Check the docs for more info.

Finally, we just need to loop through each of the generated sequences and print out the result in human readable form, using the decode() function to convert token IDs to strings.

for seq in output_sequences:
    print(tokenizer.decode(seq))

You can use the settings above to control the hyperparameters of the convolutional layer.

As you’ve probably figured out, the left side shows the input, a 20x20 handwritten digit, and the right side shows the output as it is being drawn. A convolutional layer will pass a filter over the input, selecting only a small group of pixels at a time.

Not shown: once the filter selects a window of pixels, it will compute a dot product with the pixel values and its own weights. In this demo the filter weights are assumed to be 1 for simplicity.

More on the hyperparameters

Sometimes the settings will result in an invalid convolution. This happens if the kernel extends beyond the activation map. We can write a python function to check whether it will be valid:

def valid_convolution(input_size, kernel_size, padding, stride):
  if (input_size + 2*padding - kernel_size) % stride == 0:
    return True
  else:
    return False

Padding

Increasing the padding will add extra pixels, typically zeros, around the edges of the input. This is done in order to ensure that the layer parameters will be valid. Additionally, it causes pixels which would have been on the edges to be more to the middle of the activation map.

Kernel Size

This refers to the “sliding window” which passes over the input. Choosing a larger kernel tends to increase the “smudging” which happens in the output, while a smaller kernel preserves more information for deeper layers. Sometimes kernel size is referred to as filter size.

Stride

Stride refers to the distance the sliding window will move between steps. A stride of 1 moves the window one pixel to the right, while a stride of 2 moves the window 2 pixels. A larger stride can reduce output layer size at the expense of granularity.

python -m http.server 8000

Then navigate to http://localhost:8000.

This is not meant for production environments but will get you around CORS restrictions that would come from simply opening a local file in your browser.

Let’s work backwards from where we’re trying to get to. When you’ve finished, you should be able to run the following script:

import kenlm

path = 'path/to/model.arpa'
lm = kenlm.LanguageModel(path)

sentence = "I am not superstitious but I am a little stitious"
print(model.score(sentence))

The first step will be to build KenLM. Then, we will build the ARPA file which KenLM uses to evaluate.

Building KenLM

First, clone this repository:

git clone git@github.com:kpu/kenlm.git

Now we need to build the KenLM toolkit. Run the following to build:

mkdir -p build
cd build
cmake ..
make -j 4

Now we just need to provide the model with a .arpa ngram language model file. So let’s get one.

Building an ngram language model from Wikitext-2

First, let’s clone a repository which will build an ARPA file. This repository builds the ngram file from Wikitext-2, a common dataset used in natural language processing.

git clone git@github.com:daandouwe/ngram-lm.git
cd ngram-lm
mkdir data
./get-data.sh
mkdir arpa
./main.py --order 3 --interpolate --save-arpa --name wiki-interpolate

Once that has finished, you’ll have new .arpa in the arpa directory you created. This script took the longest to run on my machine. Be patient, your computer is busy reading all of Wikipedia.

All Together Now

Now we’re finally ready to evaluate a sentence with the language model.

import kenlm

path = 'path/to/model.arpa'
lm = kenlm.LanguageModel(path)

sentence = "I am not superstitious but I am a little stitious"
print(model.score(sentence))

Which prints something like

----5---10---15---20---25---30---35---40---45---50---55---60---65---70---75---80---85---90---95--100
****************************************************************************************************
-24.47921371459961

Now if you’re interested in a bit more information about what’s going on, you can add this bit at the bottom:

words = ['<s>'] + sentence.split() + ['</s>']
for i, (prob, length, oov) in enumerate(lm.full_scores(sentence)):
  print(f'{prob} {length}: {" ".join(words[i+2-length:i+2])}')
  if oov:
    print(f'\t"{words[i+1]}" is an OOV')

for w in words:
  if not w in lm:
    print(f'"{w}" is an OOV')

Which adds this to your output:

-3.1138248443603516 2: <s> I
-1.1560251712799072 3: <s> I am
-1.1645264625549316 3: I am not
-4.912360191345215 1: superstitious
-4.504511833190918 1: but
-2.2214112281799316 2: but I
-1.1531075239181519 3: but I am
-1.2614283561706543 3: I am a
-0.9001830816268921 3: am a little
-1.2325057983398438 3: a little stitious
	"stitious" is an OOV
-2.8593297004699707 2: stitious </s>
"stitious" is an OOV

To the left of each term is the log (base 10) probability of each term occurring. For the first term, <s> I means start of sentence followed by “I”, which the model has assigned a log probability of -3.11. That’s around 0.00078. You might think it’s strange that a sentence beginning with “I” is so unlikely but we are using Wikitext-2. Wikitext-2 is Wikipedia articles. Not a lot of sentences on Wikipedia begin with “I”.

Notice that “stitious” is an OOV (out of vocabulary) term here. Clearly the language model doesn’t appreciate humor. We’ll have to tackle that next time.

Temperature is a parameter used in natural language processing models to increase or decrease the “confidence” a model has in its most likely response.

In my opinion, the most intuitive way of understanding how temperature affects model outputs is to play with it yourself. If you’re interested in the mathematical details, I’ve included them below, but I won’t be offended if you just want to play around with the slider 😃 .

Temperature (θ): 25.0

What’s going on?

Suppose we have a language model which predicts the last word in the sentence “The mouse ate the _____”. Given the previous words in the sentence and its prior training, our language model will try to fill in the blank with a reasonable final token. Suppose those raw outputs are as follows:

These outputs make sense. A mouse probably eats cheese, but mice are also known to eat cookies. A mouse probably wouldn’t eat a baguette unless it was a French mouse.

Since these are the raw outputs of the model, they won’t sum to 100. To normalize these values, we typically use softmax:

$ σ(z_i) = {e^{z_i} \over \sum_{j=0}^N e^{z_j}} $

When modulating with temperature, we introduce an additional temperature variable θ which affects the softmax distribution. A higher temperature θ “excites” previously low probability outputs. A lower temperature θ lowers the smaller outputs relative to the largest outputs. To accomplish this, we replace each z_i in the formula above with the quotient z_i/θ:

$ σ(z_i) = {e^{z_i \over θ} \over \sum_{j=0}^N e^{z_j \over θ}} $

Higher temperatures make the model more “creative” which can be useful when generating prose, for example. Lower temperatures make the model more “confident” which can be useful in applications like question answering.

TLDR: NLP metric ranging from 1 to infinity. Lower is better.

In natural language processing, perplexity is the most common metric used to measure the performance of a language model. To calculate perplexity, we use the following formula:

$ perplexity = e^z $

where

$ z = -{1 \over N} \sum_{i=0}^N ln(P_{n}) $

Typically we use base e when calculating perplexity, but this is not required. Any base will do, so sometimes the formula will use base 2 or base 10, along with logarithms to the corresponding base.

Example

Imagine that we have a language model which generates the following sequence of tokens:

<start> jack and jill went up the hill

And suppose that the conditional probabilities for each of the tokens are as follows:

For the purposes of calculating perplexity it doesn’t matter how the sequence was generated. It may be using an n-gram model or an LSTM or a transformer. All that matters is the probabilities the model assigns to each of the tokens. To calculate perplexity, we calculate the logarithm of each of the values above:

Summing the logs, we get -12.832. Since there are 8 tokens, we divide -12.832 by 8 to get -1.604. Negating that allows us to calculate the final perplexity:

$ perplexity = e^{1.604} = 4.973 $

Therefore the perplexity of this sequence is about 4.973.

http://[bucket-name].s3-website-[region].amazonaws.com

For example for “mybucket” in “us-east-1” your url will be

http://mybucket.s3-website-us-east-1.amazonaws.com

Trading agents for finance are nothing new. Previous attempts at creating automated trading systems have used statistical indicators such as moving average to determine how to act at any time. However, most of these agents focus on the action to take, opting to trade a fixed number of shares. This is not realistic for real-world trading scenarios.

This repository contains all of the code you’ll need to train and run a deep-Q agent, but if you’d like a deeper understanding of how it all works we’ve included a tutorial below.

Tutorial Overview

After following this tutorial you will have completed the following:

• Scraped four different stock indices for daily stock data trained
• Trained three deep-Q trading architectures:
  • NumQ, a simple feed-forward network which uses a single branch to determine both the action and number of shares to trade.
  • NumDReg-AD, an action-dependent network with a branch controlling the action to take and another branch controlling the number of shares to trade.
  • NumDReg-ID, an action-independent network with two branches as above, but with the number of shares entirely independent of the action taken.
• Scraped all components of the four stock indices used above and classified the components using two different methodologies.
• Trained an autoencoder, creating a neural basis for stock similarity within an index.
• Pretrained multiple models on groups of component stocks within each index.
• Trained final models on index data.
• Drawn beautiful charts displaying your results.

This tutorial assumes a basic understanding of Python and Pytorch. If you would like to brush up on Pytorch we highly recommend their tutorials.

Scraping stock indices

This repository contains a file called download_stock_data.py. Edit this file to include the names of the components for each of the following stock indices:

• Dow Jones Industrial Average (^DJI)
• S&P 500 (^GSPC)
• NASDAQ (^IXIC)
• New York Stock Exchange (^NYSE)

INDEX_NAME = 'ixic'
COMPONENTS_LIST = ['OPTT', 'NMRK', 'SCOA', ...]

Running this file will save each of the components into a directory raw/{index_name}/ with each component saved as a CSV file.

NumQ

The Jeong paper experiments with three architectures for trading. The first and simplist of these architectures is called NumQ which uses a single branch of fully-connected layers to determine both the action to take and the ratios for those actions. Its structure is shown below:

This can be succinctly represented in pytorch as a series of torch.nn.Linear layers:

class NumQModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(in_features=200, out_features=200, bias=True)
        self.fc2 = nn.Linear(in_features=200, out_features=100, bias=True)
        self.fc3 = nn.Linear(in_features=100, out_features=50, bias=True)
        self.fc_q = nn.Linear(in_features=50, out_features=3, bias=True)

    def forward(self, x: Tensor) -> Tuple[Tensor, Tensor]:
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)

        q = self.fc_q(F.relu(x))
        r = F.softmax(self.fc_q(torch.sigmoid(x)))

        return q, r

Of note in this model are the return values of the forward() function. Rather than simply returning one result of a forward-pass through the network, the NumQ network will return two values: a list of Q values associated with the input, and the ratios for each of the three actions (BUY, HOLD, SELL). When in use, the Q values determine which action will be taken, while the ratios determine with how many shares the action will be executed with.

NumDReg-AD

The second model is a bit more complex. It contains a branch for determining the Q values associated with the input, as well as a separate branch for determining the numbers of shares to trade.

This architecture is quite a bit more involved than NumQ. In fact, this increase in complexity motivates a reduction in each of the fully-connected layers in the network from those in NumQ: 100, 50, 20 versus the 200, 100, 50 used in NumQ. In spite of this adjustment, this and the next model, NumDReg-ID, need to be trained using a three-step training process. The training process will be discussed later on.

class NumDRegModel(nn.Module):
    def __init__(self, method):
        super().__init__()

        # Set method
        self.method = method

        # Training step
        self.step = 0

        # root
        self.fc1 = nn.Linear(in_features=200, out_features=100, bias=True)

        # action branch
        self.fc2_act = nn.Linear(in_features=100, out_features=50, bias=True)
        self.fc3_act = nn.Linear(in_features=50, out_features=20, bias=True)
        self.fc_q = nn.Linear(in_features=20, out_features=3, bias=True)

        # number branch
        self.fc2_num = nn.Linear(in_features=100, out_features=50, bias=True)
        self.fc3_num = nn.Linear(in_features=50, out_features=20, bias=True)
        self.fc_r = nn.Linear(in_features=20, out_features=(3 if self.method == NUMDREG_AD else 1), bias=True)

    def forward(self, x: Tensor) -> Tuple[Tensor, Tensor]:
        # Root
        x = F.relu(self.fc1(x))

        # Action branch
        x_act = F.relu(self.fc2_act(x))
        x_act = F.relu(self.fc3_act(x_act))
        q = self.fc_q(x_act)

        if self.step == 1:
            # Number branch based on q values
            r = F.softmax(self.fc_q(torch.sigmoid(x_act)))
        else:
            # Number branch
            x_num = F.relu(self.fc2_num(x))
            x_num = torch.sigmoid(self.fc3_num(x_num))
            # Output layer depends on method
            if self.method == NUMDREG_ID:
                r = torch.sigmoid(self.fc_r(x_num))
            else:
                r = F.softmax(self.fc_r(x_num))

        return q, r

    def set_step(self, s):
        self.step = s

NumDReg-ID

This is the third and final paper introduced in the paper. It contains an action-independent branch which specifies the number of shares to trade. Its architecture is identical to NumDReg-AD except for the output layer in the number branch which has a size of one. The activation function of the number branch was not specified in the paper, so we opted to use the sigmoid function.

(Sigmoid has the convenient property of being bounded between 0 and 1. This avoids the awkward circumstance where the action branch predicts BUY while the number branch predicts a negative number, resulting in a SELL action.)

Reinforcement learning

At each timestep t, the agent observes a state of 200 price differences, chooses an action (either BUY, HOLD, or SELL) with a number of shares, and receives a reward equal to the normalized profit.

Training process

Reinforcement learning agents are trained over a number of episodes, during which they observe states, take actions, receive rewards, and observe the next state. These (state, action, reward, next_state) transitions are stored in a memory buffer which the agent then uses to optimize its neural network.

The learning process for deep Q networks is a bit different from normal Q learning models. Deep Q models typically contain two neural networks working in tandem: a policy network which evaluates a given state, and a target network which is periodically updated with the weights from the policy net. This periodic update pattern helps to maintain stability while training.

Our training logic defines an episode as one chronological pass through the training data. This detail is not specified in the paper, but one pass over the data makes logical sense in this context. We used a TradingEnvironment class to track information during training, which has the added benefit of making the code more readable. The details of TradingEnvironment will be discussed later on.

At the beginning of training, the policy network and target network are initialized. After this, we begin to iterate over a number of episodes.

During an episode, the agent is provided with a state, which you will recall is a number of price differences. This state is fed into the policy network, which will calculate Q values and ratios. We then calculate the reward for each of the three actions. The (state, action, reward, next_state) transitions for each of the actions are stored in a memory buffer. Next, the model will undergo an optimization step.

During optimization, batch transitions are retrieved from the memory buffer. Then the loss is computed as the difference between actual and expected Q values and backpropagated through the policy net.

After an episode concludes, we do an update of the target network with the policy network and reset the environment to begin serving states from the beginning of the episode again.

3 step training (NumDReg-AD/ID)

1. Train the shared root and action branch only. The number branch is frozen.
2. Train the number branch from the action branch’s Q-values. The action branch is frozen.
3. Train the entire network end-to-end.

Each step runs for the same number of episodes.

Component stock relationships

Because of their large number of weights, deep neural networks have the tendency to overfit their training data if there isn’t enough of it. To counteract this behavior, we reduced the number of weights and changed the training process into a 3-step procedure. One other thing which can help to prevent overfitting is to train with more data. This is why we pretrain using component stocks and finish training on the index stock.

But first we need to choose which components to pretrain with. Training with all of them will be too time-consuming. Instead, we will create 6 groups of stocks based on 2 different measures of the components: Pearson correlation and the mean squared predicting error from an autoencoder. The autoencoder implementation will be discussed in the next section. We will measure each component stock and choose 6 groups according to the following:

So for each stock we will measure its correlation with the index and measure its autoencoder MSE. One group of stocks will be made up of the highest 2n components when measured by correlation with the index. (Here n is chosen in proportion to the number of components in the index.) Another group will be made up of the 2n stocks which have the lowest correlation with the index. And so forth.

Creating these groups allows our models to be pretrained on stocks which are proxies to the index, or in the case of the “low” group, on stocks which might help the agent generalize later on.

Autoencoding component stocks

In order to calculate the mean squared error of the component stocks, we need to train an autoencoder which will predict a series of stock prices for each component in an index. That is, the input size of the network will be MxN, where M is the number of components in the index and N is the number of days in the time series.

We will train an autoencoder such that X=Y, where X is the input and Y is the output. The architecture of the autoencoder is very simple, having only 2 hidden layers with 5 units each. These small hidden layers force the model to encode the most essential information of its inputs into a small latent space. All extraneous information not represented in the latent space is discarded. Each of the inputs x_i will be encoder with the autoencoder as y_i, and it is against these that mean squared error is measured.

Training an autoencoder is fairly simple if you understand the basics of neural networks. Rather than training with (input, target) pairs, since we only care that input=input, we will train on (input, input) pairs.

Transfer learning

Financial datasets are small, so we pretrain on grouped component stocks, then transfer those weights to train on the actual index. Stock groupings are chosen by two measures — Pearson correlation with the index and autoencoder reconstruction MSE — each producing a “high 2n,” “low 2n,” and “high n + low n” group.

When the max Q-value for an action is within 0.0002 of the min, the market signal is too weak to act on, so the agent defaults to HOLD.

Finance Environment

For convenience, we include an environment for training which encapsulates many variables that would otherwise need to be tracked in the training loop. Instead, we use the FinanceEnvironment to store these variables and provide them as needed to the agent during training.

The FinanceEnvironment class exposes only a few necessary methods to the training loop. For example, step() returns state and done. The first of these is the difference in stock prices for the 200 days prior to the current day. The second indicates whether the episode is done. Executing step() updates several internal variables used in profit and reward calculations.

The other important method of the environment is update_replay_memory(). This method adds a (state, action, rewards_all_actions, next_state) transition to the replay memory. This will be sampled later when the model is optimized. Because the environment stores and updates most of these variables internally, they do not clutter up the training loop.

Putting it all together

You can see our implementation in the github repo.

python main.py --train --architecture numdreg_ad --index ixic

See python main.py --help for all options: architecture selection, pretraining, tensorboard logging, and checkpointing.

Selected Results

NumQ and NumDReg-ID generally outperformed NumDReg-AD across indices, with the best returns on ^GSPC (S&P 500). Final cumulative profits fell short of the paper’s reported figures — likely due to unreported hyperparameters, a shorter training window, and the inherent difficulty of equity prediction. Full result tables and Sharpe ratio comparisons are in the repo.

1. These articles are mainly a reference for me. I find myself searching the same things over and over, looking for the purple link, scrolling through the article, then copy & pasting code. I’d rather not go through the hassle. These articles aim to solve that problem.
2. I aim to keep the answers above the fold. I don’t want to have to scroll down to find the answer. I almost never read the surrounding prose when I am in “coding mode”.
3. I don’t have ads or popups on my blog. I will never ask people to sign up for a newsletter or login to read more. I also don’t use pictures unless there’s a good reason. The “thinking AI robot stock photo” industry is definitely a bubble.
4. Writing these things out explicitly helps me to remember them. Paradoxically, this may make these how-to pages less useful to me, but maybe someone else will find them useful.

These quick reference articles don’t explain much because I don’t need an explanation of what is going on. I just need a chunk of working code. There are other websites which have far more comprehensive guides covering how to do things. They cover all of the fundamentals of how things are done. But I don’t need that, I just want a 30 second reference with a working chunk of code.

Articles

• Python: Formatting a string
• Python: Read & Write Json
• Python: Serve an HTML File
• Managing Python Environments
• Colab: Connect to Google Drive
• Connect Jupyter to Remote
• S3 Bucket Url
• How to Zip and Unzip a tar.gz File

“God bless JSON!” ~ a soon to be famous programmer

import json

data = {'a': 1, 'b':'hello', 'c':False}
filename = 'awesome_data.json'

# write data to file
with open(filename, 'w') as f:
  json.dump(data, f)


# read json from file
with open(filename, 'r') as f:
  data = json.load(f)


print(data)
# prints {'a': 1, 'b':'hello', 'c':False}

So you want to build an autoencoder? Great! This article will demonstrate how to build an autoencoder and use it to measure stock prices against an index. This technique is described in more technical terms here.

Once we’ve trained the autoencoder, we can use it to measure how well each component follows the other members of the index. This can be useful for finding deeper insights into an index, and doesn’t require a priori knowledge of the index price or the weighting of its components. Note, this is only one metric which one could use to determine how well one member of the group follows the group overall. Another might be Pearson Correlation.

Github repository

To follow along with the code in this tutorial, please download the corresponding repository on Github:

git clone git@github.com:lukesalamone/stock-data-autoencoder.git
cd stock-data-autoencoder
pip install -r requirements.txt

What is an autoencoder?

An autoencoder is a neural network which encodes information back to itself. The structure of the network is such that the input layers (the “encoder”) will be large compared to the hidden layers (the “code”), forcing the network to compress information inside its hidden layers.

The idea of our autoencoder is that we would like to encode stock price information back to itself while discarding trends that aren’t important. To do this, we will feed our network stock price data and ask the network to return those prices to us as outputs. Component stocks which are important to the index will be preserved well and thus be highly accurate, while components which are less important will not be well-preserved. We will measure the performance of the network on each component using mean squared error.

The model

We will use an autoencoder with a number of inputs and outputs equal to the number of component stocks in our index. For this exercise, we will use the S&P 500 index which contains 505 components. This means our input and output size will be 505. We will also use a hidden layer with 5 units.

class StonksNet(nn.Module):
 def __init__(self, size):
   super().__init__()
     self.fc1 = nn.Linear(in_features=size, out_features=5)
     self.out = nn.Linear(in_features=5, out_features=size)

 def forward(self, x: Tensor) -> Tensor:
   x = F.relu(self.fc1(x))
   x = F.relu(self.out(x))
   return x

The data

We will use daily stock prices downloaded using yfinance. This data is readily available online and I recommend downloading it for yourself. We will use data between January 1, 1991 to January 1, 2021 (30 years of data).

To download the S&P 500 stock data please run gather_stocks.py from the project directory:

python gather_stocks.py

This will download all 505 components into the stock_data directory. Data will also be cleaned such that each component has the same number of days, which will be important when feeding it into the model.

Training the model

The model itself is a simple feed-forward neural network. As such, we use a standard training loop to train the model. We don’t expect the loss to ever fall to zero during training since it is impossible for the network to perfectly encode and decode so many inputs into so few hidden code units. Some information will inevitably be lost. In my training, validation losses bottomed out at around 4000, but yours may be different depending on the initialization of your autoencoder.

Ranking components

Finally we’re ready to rank the components of the S&P 500 for “closeness”. After running python train_model.py you will see the best and worst components as scored by the autoencoder. Here were my results, yours may be different.

best 5 results:
DRE: 16.66
LNT: 37.27
MU: 38.88
HOLX: 43.18
CERN: 47.46

worst 5 results:
HUM: 105244.19
SHW: 108542.73
LMT: 113654.48
C: 357073.88
NVR: 10955169.00

Future research

Upon inspection, it appears that better results might be achieved if we normalize the stock data before training. It appears that stocks with higher prices and higher volatility tended to perform worse than those with tight price ranges. In a way this is expected, since the autoencoder will naturally have a harder time modeling large values with a limited set of hidden units. However, normalizing the prices into similar ranges might be an interesting exercise to see if we can squeeze even more out of the model.

name = 'Luke'
food = 'pizza'

# old style
"My name is %s and I like %s." % (name, food)

# str.format()
"My name is {0} and I like {1}.".format(name, food)

# f-strings
f"My name is {name} and I like {food}."

The following is a transcript of the above video

In this paper, the authors present a novel neural network architecture to enable audio search via sounds humans are able to make, for example humming and whistling. This is an important capability when searching through audio for a specific sound.

Motivation

Imagine you have hundreds of unlabeled sound effects on your computer, and you are looking for a specific one. It could be very tedious to listen to every single one until you can find the right sound. Even if the sounds do have some kind of word labels, it could be hard to pinpoint exactly which words to search for. A lot of sounds don’t exactly lend themselves to text descriptors, so finding the right sound can be difficult with a text search.

This paper contains three main contributions:

First, it introduces a new neural network architecture for matching imitated sounds with a sound corpus, the semi-siamese convolutional neural network.

Second, the researchers built a second architecture which utilizes transfer learning from other audio tasks in an attempt for better performance.

Third, the researchers visualized and sonified input patterns which excited the neurons in different layers.

Both neural networks outperform the state of the art systems as we will see later on.

Data

To train the siamese model, the authors used a dataset called VocalSketch, which contains sounds from 4 broad categories: acoustic instruments, commercial synthesizers, everyday, and single synthesizer notes. The dataset also contains a number of human vocal imitations of each of the sounds. For each of the 4 categories in VocalSketch, the researchers selected half of the sounds in each category along with corresponding vocal imitations as the training and validation set, and the other half for their test set.

Since each of the categories other than Everyday contained 40 sounds, 20 sounds would be for training and 20 for test and validation. The everyday category contained 120 sounds, so that category had 60 in training and 60 in test. Each sound in the dataset has 10 corresponding human vocal imitations. In the training set, 7 of them for each sound were selected for training, and the remaining 3 were used for validation. Overall, the researchers used this dataset to create 840 positive and 840 negative pairs for training, and 360 positive and negative pairs for validation.

I thought it was interesting that the researchers opted not to use balanced categories of 20 sounds for each category. There isn’t a comment in the paper about the reasoning behind this but it may be due to the difficulty of categorizing this class of sounds.

The transfer learning model required pretraining the two towers on two additional datasets before training the full network on VocalSketch.

For the vocal imitation tower, they used a dataset called VoxForge. From this dataset the researchers selected 8 thousand samples for each of 7 different languages. They used a 70 - 30 split for training and testing, and achieved a 69.8% accuracy, which seems pretty good for a 7 class classification task.

For the second transfer learning tower, the environmental sound classification tower, they used a dataset called UrbanSound8k. This dataset contains 8732 sound samples in 10 different classes, things like car horns, jackhammers, and street music. The researchers used 10 fold cross validation when pretraining on this dataset and achieved a 70.2% accuracy over the dataset.

Note that for pre-training the two towers, the researchers used a slightly modified neural network architecture, appending two fully connected layers to categorize the results into the necessary number of classes.

So how are sounds fed into the neural networks? You’re probably familiar with the way that convolutional neural networks work with images. Typically the input for each image is of a shape width by height by color depth. Creating an input when working with sound is similar. The width of the audio file in this case is time, and the height is the frequencies during that time step. This is similar to the output of a spectrogram.

In order to be fed into the network, audio must first undergo a preprocessing step. The specific preprocessing involved varied between the networks, but generally involves downsampling the audio and splitting it by frequency band. This resulted in an input which resembles a spectrogram image.

Methodology

The heart of this problem is an architecture the researchers dubbed siamese style or semi siamese neural networks. A true siamese neural network consists of two identical neural network towers with the same weights, and is used similar to the way that hashing is used to match similar inputs. A siamese style network is similar to a siamese neural network, but may have different weights in one of its towers.

The first network the researchers called IMINET, which included a true siamese neural network as one of its configurations. This network consisted of two convolutional network towers, a concatenation step, and a fully-connected network with three more layers to compute a similarity score between 0 and 1. The convolutional neural networks in IMINET had the same structure, even though in some configurations their weights were not the same. They consisted of four convolutional layers with pool layers following convolutional layers 1 and 2.

The researchers detailed the parameters they used for each of the convolutional layers, specifically the number of filters and the receptive field size. Each of the filters in a layer learns to detect various characteristics of the input, and the receptive field is the part of the input that the filter is able to see. The max pooling layers output the maximum value among all of their inputs, reducing the number of inputs for the next layer in the network.

The researchers experimented with three different configurations for the convolutional network towers: a tied configuration, where both towers would share the exact same weights and biases; an untied configuration, where the two towers were required to share no weights and biases at all; and a partially tied configuration, where weights and biases were shared for convolutional layers 3 and 4 but not 1 and 2. Because the untied and partially tied configurations are not truly siamese neural networks, the researchers called these configurations semi siamese.

After both inputs pass through the convolutional towers, they are concatenated together into one input vector and fed into a fully connected network with three layers. This network’s job is to compute a similarity score between 0 and 1 for the two inputs. Both of the first two layers used ReLU activation. The final layer was a single neuron which used sigmoid activation to squash its input into an output between 0 and 1.

This architecture achieved state of the art results in sound retrieval which I will discuss in a moment along with the rest of the findings of this paper.

The second neural network developed by researchers was called TL-IMINET. This network was very similar to the first network, but this time the researchers tried using transfer learning to achieve better performance. The researchers hypothesized that the vocal imitation task shares many characteristics with language identification, and that sound recognition shares many similarities with environmental sound classification. This would allow networks which were pre-trained on these tasks to require only fine tuning to be adapted to this task.

The network architectures for the language identification and environmental classification tasks were slightly different from those used in IMINET and are shown here. Note that the towers are also different in architecture from each other.

The researchers also experimented with fusion strategies between different models. For IMINET, the similarity likelihoods for all three configurations were multiplied together to achieve a combined score. They also experimented with combining IMINET with the previous state of the art model. Since that model computes a cosine distance between the input sound and a candidate sound, this output was converted into a likelihood using softmax, and that softmax was multiplied by the output of IMINET. The transfer learning model TL-IMINET was also combined with the state of the art model in a similar way by computing the softmax and multiplying by the output of TL-IMINET. These fusion strategies ended up improving the performance of each of the models quite a bit.

Experiments and Findings

To measure the performance of these networks, recall that the output of the networks was a number between 0 and 1 indicating how similar the network believed the two inputs were to each other. For example, two very similar inputs might have a similarity rating of 0.9, while two dissimilar inputs might have a similarity rating of 0.1. After gathering the similarity of the human vocalization sound to each of its potential matches, the matches were ranked according to their similarity score.

The authors then used a metric called mean reciprocal rank, which is a number between 0 and 1 indicating how well the algorithm ranked the sounds. For example, a mean reciprocal rank of 0.5 suggests that the target sound is ranked second among all possible matches on average.

Here are the performances of the various network configurations when measured by mean reciprocal rank. The researchers highlighted several insights which could be drawn from their results.

First, it seems that tied configurations performed the best among all configurations of IMINET. This runs contrary to the researchers’ expectations that untied configurations would outperform tied configurations.

Second, the tied configuration outperformed the previous state of the art benchmark in two categories of sounds: commercial synthesizers and everyday. It performed worse than the state of the art for acoustic instruments, and was about the same performance for the single synthesizer category.

Third, IMINET achieved even better performance in most categories by using an ensemble of different configurations.

Fourth, even without pretraining, the TL-IMINET model performed better than the untied configuration of IMINET for all categories except commercial synthesizers. This is interesting because the only difference between these two models is the network structure of the convolutional towers.

And finally, the pre-trained TL-IMINET model outperformed the previous state of the art model by quite a bit in all categories, but the best performing configuration overall was TL-IMINET fused with the previous state of the art model.

One of the most interesting experiments was the visualization of the input patterns which activate neurons in each of the layers. This was done by performing gradient ascent of the neuron activation with respect to the input from a random initialization state. Visualizing the convolutional layers showed that the first layer tends to learn local features like edges, intermediate layer neurons learn more complex features like texture and direction, and the deepest layer recognizes concentrations of frequency ranges. The visualizations also helped to confirm that pretraining indeed helped the networks to learn more detail. The patterns from the pretrained vocal imitation tower were sharper than those in the naive IMINET towers.

Takeaways

There are a few key takeaways from this research. The first is that the transfer learners had much better performance than the naive non-transfer learner, as evidenced by the fact that TL-IMINET performed better than IMINET for most categories even though neither model was pretrained. The research also showed that ensemble methods can outperform any single model on its own. IMINET performed better when used in combination with its different configurations, and combining it with the state of the art model performed better than either model on its own. Finally, visualizing the inputs can help to confirm that the network is learning the correct things, and helps to provide insights as to what types of sound properties are important.

Are you studying machine learning and want to know more about Gaussian Mixture Models? You’ve come to the right place. I have found other online resources to be difficult to approach and/or lacking crucial details. Here I will try to explain GMMs in plain language.

Overview

Gaussian Mixture Models are used to group data points into different clusters. In this way, it is similar to K-means clustering, where each cluster is specified by a point at its center. In a way, GMMs are a generalization of K-means, where we relax the assumption that the clusters are circular (or spherical, or n-dimensional sphere, depending on how many dimensions your features have). Gaussians are specified by three parameters. Just like how the K-means algorithm seeks a point which best “fits” your data points, the GMM algorithm will seek a gaussian which best fits them.

Any GMM model you will be implementing will probably have a fit and a predict method. The goal of the fit method is to find a number of gaussians to cluster your data points. The predict method takes a new point (or a list of new points) and predicts which gaussian is closest to it. So you may have a skeleton class like so:

class GMM():
    def __init__(self, n_clusters, covariance_type):
      pass

    def fit(self, features):
      pass

    def predict(self, features):
      pass

The fit() method

This method has two major sub-parts: the expectation step, and the Maximization step. These two parts will be run inside of an update loop which typically terminates after a set number of iterations or the algorithm has stopped improving.

I personally find that examples are the best way to understand many concepts, so in this article I will use a dataset with 4 features, each of the features having 3 dimensions:

Expectation Step

The purpose of this step is to calculate for each of our features, the probability of that feature belonging to cluster k. In this example we have 4 clusters and 4 data points. So we need to fill in the following table:

Maximization Step

In this step we need to update three variables for each of the clusters: the weight, the mean, and the diagonal covariance. In this article I will use w_j to represent the weight of the j^th gaussian, μ_j to represent the j^th mean, and σ_j to represent the j^th variance.

Calculating w_j

The first variable will be r_j, which is the sum of gammas corresponding to a given cluster k_j. In other words,

$ r_{j} = \sum_{i=0}^3 γ_{ij} $

Basically, r_j is the sum of the gammas in the j^th column. We will use this r_j to calculate a w_j for each cluster, the “mixing weight” of the gaussian. Divide by the number of features (in this case 4):

$ w_{j} = {r_{j} \over 4} $

Calculating μ_j

Next we calculate the means μ_j of the gaussians i.e. the point in n-dimensional space over which the (n+1)-dimensional bell curve is centered. Because in our example each feature x_i contains 3 dimensions (x₀ = [d₀₀, d₀₁, d₀₂]), the means will also need to have 3 dimensions each.

μ_j is calculated as follows:

$$ μ_{j} = {1 \over r_{j} } \sum_{i=0}^3 γ_{ji} x_{i} $$

Let’s work through the first μ. For simplicity we will multiply by r_j for now. For the first mean μ₀:

$$ r_j μ_0 = γ_{00} x_{0} + γ_{01} x_{1} + γ_{02} x_{2} + γ_{03} x_{3} $$

$$ = γ_{00} \begin{bmatrix} d_{00} \\ d_{01} \\ d_{02} \end{bmatrix} + γ_{01} \begin{bmatrix} d_{10} \\ d_{11} \\ d_{12} \end{bmatrix} + γ_{02} \begin{bmatrix} d_{20} \\ d_{21} \\ d_{22} \end{bmatrix} + γ_{03} \begin{bmatrix} d_{30} \\ d_{31} \\ d_{32} \end{bmatrix}$$

So we are essentially multiplying the gammas associated with a gaussian by all of the features we have. This matrix multiplication can be rolled up into

$$ = \begin{bmatrix} d_{00} & d_{10} & d_{20} & d_{30} \\ d_{01} & d_{11} & d_{21} & d_{31} \\ d_{02} & d_{12} & d_{22} & d_{32} \end{bmatrix} \begin{bmatrix} γ_{00} \\ γ_{01} \\ γ_{02} \\ γ_{03} \\ \end{bmatrix} = \begin{bmatrix} μ_{00} \\ μ_{01} \\ μ_{02} \end{bmatrix} $$

Don’t forget to divide by r_j again! It makes sense that each μ should have 3 dimensions, since our features also each have 3 dimensions.

Calculating σ_j

Technically this is σ² if you’re here from statistics, but since we never take the square root it doesn’t matter for our purposes.

All together now. E step with code.

The full equation for the expectation step is

$$ \gamma_{ij} = \frac{ w_j \, \mathcal{N}\!\left(x_i \mid \mu_j, \Sigma_j\right) }{ \sum_{k=0}^{K-1} w_k \, \mathcal{N}\!\left(x_i \mid \mu_k, \Sigma_k\right) } $$

Alright, enough math. As fun as it is for you to read, it’s even more fun to write in markdown.

All of the math above for calculating w, μ and σ can be written succinctly in python. Here I am using numpy to take advantage of some pretty slick matrix operations they have available.

"""
Fit GMM to the given data using self.n_clusters number of gaussians.
Features may be multi-dimensional.

Args:
  features: numpy array containing inputs of size
    (n_samples, n_dimensions)
  expectations: numpy array containing gammas from E step. Size is
    (n_features, n_clusters)
Returns:
  means: updated means
  covariances: updated covariances
  mixing_weights: updated mixing weights
"""

def maximization_step(self, features, expectations):
  # initialize lists to hold our parameters
  mixing_weights, means, covariances = [], [], []

  for cluster in range(self.n_clusters):
    # gammas will hold a column from the expectations matrix
    gammas = expectations[:,cluster]

    # r is just the sum of a column
    r = np.sum(gammas)

    # calculate mixing weight
    w = r / len(features)
    mixing_weights.append(w)

    # calculate means
    # the @ symbol is a slick shortcut for matrix multiplication
    mu = (gammas @ features) / r
    means.append(mu)

    # calculate covariances
    diff_sq = (features-mu) ** 2
    cov = (gammas @ diff_sq) / r
    covariances.append(cov)

  return means, covariances, mixing_weights

The predict() step

Once the model is trained, we have learned:

• Mixing weights w_j
• Means μ_j
• Covariances Σ_j

To classify a new data point x we compute its responsibilities using the same formula from the E-step:

$$ \gamma_j(x) = \frac{ w_j \, \mathcal{N}\!\left(x \mid \mu_j, \Sigma_j\right) }{ \sum_{k=0}^{K-1} w_k \, \mathcal{N}\!\left(x \mid \mu_k, \Sigma_k\right) } $$

This yields a probability distribution over all clusters. If we only want the most likely cluster,

$$ \text{cluster}(x) = \arg\max_{j} \, \gamma_j(x) $$

def predict_probs(self, features):
    gamma, _ = self._e_step(features)
    return gamma

def predict(self, features):
    return np.argmax(self.predict_probs(features), axis=1)

Installing pyenv

# install pyenv
curl https://pyenv.run | bash

# check pyenv install location
which pyenv

Install another python version

# see a list of available python versions
pyenv install --list

# check installed python versions
pyenv versions

# installs python 3.7.5
pyenv install 3.7.5

Switch python versions

# use python 3.7.5 everywhere on your machine
pyenv global 3.7.5

# use python 3.7.5 in current directory
pyenv local 3.7.5

# use python 3.7.5 in current shell session
pyenv shell 3.7.5

Algorithm

The basic K-means algorithm is fairly simple and has two steps, repeated until convergence:

1. assign points to cluster corresponding to closest centroid
2. update centroid locations to the mean of all points assigned to the associated cluster

The algorithm converges when the centroids stop moving, i.e. no points can be switched between clusters to a closer centroid.

def kmeans(points, k):

  random.shuffle(points)
  centroids = points[:k]
  done = False

  while not done:
    clusters = clear_clusters()

    # assign each point to closest cluster
    for point in points:
      cluster_idx = closest_cluster(centroids, point)
      clusters[cluster_idx].append(point)

    # recalculate centroids
    new_centroids = update_centroids(clusters)
    done = new_centroids == centroids
  return clusters

It seems like a simple enough question. Which word should you choose so that it takes your opponent the most guesses to discover it? Should you choose a long word to use up your opponent’s guesses? Or perhaps a short word with obscure letters? In this document I look into this question. But first, a bit of background.

If you’re not familiar with the rules of hangman, it is a guessing game played between two people. Player A chooses a secret word, and tells player B the length of the secret word. Player B guesses letters which she thinks might be in the word. If she chooses a correct letter, player A reveals the locations of each instance of the guessed letter. However, if player B guesses an incorrect letter, this counts as a “strike” against her. After an agreed-upon number of strikes, player B loses.

An algorithmic approach

A few years ago, I created a “hangman solver” for the popular paper and pencil game. This game assessed each game analytically, to determine a list of possible words given the clues available. The algorithm works as follows: at the beginning of the game, we know the length of the secret word, which narrows our dictionary considerably. Then, for each letter in the alphabet, count the number of words available which contain that letter.

Suppose our dictionary consisted of a random list of 50 four-letter words as follows:

["pull", "dipt", "dorp", "poky", "jism", "cues", "hood", "drag",
"inky", "mhos", "kerf", "jess", "mete", "lues", "wipe", "kane",
"tiro", "keys", "jape", "lime", "sees", "sass", "demo", "ilia",
"mink", "dips", "hove", "jees", "that", "pops", "isle", "teas",
"dens", "dogy", "pink", "sizy", "cole", "pact", "thaw", "lead",
"mile", "dodo", "litu", "scup", "colt", "soma", "seat", "dewy",
"pits", "mojo"]

This would result in letter counts as follows:

In this case, the letter E is found in 26 words, the most of any letter. Therefore our algorithm should pick E since it is the safest guess. Once we guess a letter correctly, this gives important positional information which can filter the word list even further.

This process is repeated, each time picking the most likely letter, given the constraints we know. If we know some of the letters in the secret word, we can eliminate any words that don’t have those letters in those positions. If we have guessed a letter incorrectly, we know that letter isn’t in the secret word and can eliminate all words which have that letter. You can see a python implementation of this algorithm here.

For this experiment I used the Zynga dictionary, the same dictionary used in the game Words With Friends. The dictionary contains 173,000 words. It does not include any proper nouns or profanities.

Live experiment

Below you can experiment with this algorithm among 4 letter words. Enter a secret word, and the steps used to uncover the secret word will be displayed below. (Need inspiration? Try comparing jazz to rock or blue to grey.)

Notice how some words have much higher difficulty than others? This is due to the fact that some words have many “siblings” which differ by only one letter. For example the pattern “.ays” could match the letters c, d, f, h, j, k, l, n, or r. Knowing the last 3 letters gives us no indication of which of these nine letters it will be.

One apparent shortcoming of the above calculation is that it assumes that letters with equal probability will be picked in alphabetical order, and therefore letters last in the alphabet will be picked last. Although this has the benefit of creating a deterministic algorithm which will always return the same result for the same word, in real life we don’t know in which order people will pick letters. In actuality the words days, jays and rays have equal difficulty each is equally likely to be the secret word.

Preliminary results

With this preliminary caveat in mind, we can still calculate the difficulty of every word in the dictionary. If we assume that in a situation where multiple letters are equally probable, our opponent will break the tie using alphabetical ordering, which hangman words are the hardest to guess? Here are the top 17:

Previous research

I should note that previous research by Jon McLoone in 2010 has explored the same topic, although he used slightly different methodology and a smaller 90,000 word dictionary. His algorithm was not deterministic, and so does not always pick the most frequent letter available. Instead, his algorithm picks a letter with a probability proportional to the frequency with which it occurs in candidate words. For example, if we refer to the letter frequencies of the 50 word 4-letter dictionary above, j appears in just 1 word, while e appears in 26 of them. In this case, since the sum of the numbers in the frequency table is 188, McLoone’s algorithm would pick z in 1 out of every 188 first guesses, and e in 26 of them. Although it might seem that this strategy is not optimal, it does avoid the deterministic results shown above.

Additionally, McLoone chose to remain faithful to the logic of the hangman game, opting to end the games after a given number of mistakes, and recording the probability that a given word was not discovered after the game ended. So an 8-game means that the game was ended after 8 mistakes, and a 13-game after 13 mistakes. Using this methodology, he found the hardest hangman words were as follows:

Future research

Now, I think that we can improve upon our results a bit. Rather than calculating difficulty deterministically, we can instead randomize the ordering which letters will be picked in. This should dramatically reduce some of the outliers from above, bringing “rays” down from a difficulty of 14 to something more reasonable.

Such a stochastic calculation will require simulating millions of games. For my 173,000 word dictionary, simply simulating 10 games would involve 1.7 million games. Fortunately, this operation is highly parallelizable. It should be possible to split the dictionary into 100 or even 1000 pieces and derive the results for each piece simultaneously.

A simulation of this algorithm is shown below, with a graph of the average number of mistakes the new randomized algorithm incurs before discovering your secret word. The simulation is set up to play 500 rounds of games, and the final average is displayed at the top.

A Monte Carlo simulation is a method of estimating events or quantities which are difficult or computationally infeasible to derive a closed-form solution to. The value of the mathematical constant Pi is a good example of this: although it is possible to calculate the exact value of Pi, a good estimate is easily demonstrated with just a few lines of code.

Monte Carlo simulations work when the input can be drawn from a random probability distribution, and the outcome can be derived deterministically from the input. In the case of calculating Pi, this can be modeled geometrically. The random distribution is all points within the square, and the outcome is whether a selected point lies within the circle inside of the square. We know that for a square circumscribed about a circle,

$$ Area_{circle} = \pi r^2, Area_{square} = (2r)^2 = 4r^2 $$

$$ \frac{Area_{circle}}{Area_{square}} = \frac{\pi r^2}{4r^2} = \frac{\pi}{4} $$

$$ 4(\frac{Area_{circle}}{Area_{square}}) = \pi $$

If we notice that the probability that a randomly placed dot will fall within the circle is the same as the ratio of their areas (i.e. the circle takes up about 78% of the area of the square, so a random dot has about a 78% chance of landing inside the circle), then multiplying that probability by 4 gives Pi. By placing dots randomly, we play out that probability in real-time.

As to whether a given dot lies within the circle, we simply use the Pythagorean theorem to calculate its distance from the origin:

$$ \sqrt{ x^2 + y^2 } < 1 $$

In other words, all dots greater than 1 unit from the origin are outside the circle.

Below is a simulation of the derivation of the value of Pi. Click “start simulation” to see for yourself. Warning: this simulation may become slow once many dots are drawn on the screen.

Next, I was interested in the way that this simulation would play out once the number of points became large. Intuitively, this should approach Pi, but doing so requires that the random numbers generated by browsers be evenly distributed.

In February 2020 I decided to code up Gomoku2049. The game is a demonstration of MiniMax, which is an algorithm for finding the move which minimizes the opponent’s best moves. This article is an overview of the game’s technical highlights.

Click here to try out the game!

Minimax with alpha-beta pruning

The Minimax algorithm represents every game as a tree of moves, with the current game position at the root of the tree. The algorithm is recursive with exponential time complexity and can have a very high branching factor: after the first move there are 225–1=224 possible moves. Because it is not feasible to evaluate all possible games to completion, Minimax calculation is usually limited to a fixed depth, after which the algorithm evaluates terminal leaf nodes using the gameover function and the static evaluator.

After each human move (known as “plies”), Minimax assigns a score to each of the possible reply moves. By convention, the AI will score favorable moves with a positive score, and unfavorable moves with a negative score. The move corresponding to the highest score is then selected. In other words, the AI is called the “maximizer”. Likewise, the human is known as the minimizer. To determine the score of each possible move, the minimax algorithm will recursively either maximize or minimize the possible moves available. After a given depth, the evaluation will stop, and return either an infinite value (+∞ for an AI win, -∞ for human win) or a finite evaluation of the state of the board. This static evaluation can be rather expensive, but luckily even a rough approximation is effective.

In practice, in addition to a depth limitation, this minimax algorithm also reduces the branching factor by limiting the squares which will be evaluated to those which are adjacent to squares which have been played. Given the fact that a disconnected “island” square cannot immediately lead to a win, this seems to be a reasonable simplification. At the leaf nodes of the tree, either the game is over (the human has won or the computer has won) or the board needs to be evaluated with regards to who is winning.

Alpha-beta pruning

Alpha-beta pruning is an improvement on the minimax algorithm, reducing the number of branches and leaf nodes which need to be evaluated. This is achieved by “pruning” unnecessary branches, ignoring them because the parent minimizer/maximizer would never choose it. For a maximizer (whose parent is a minimizer), this will occur if the parent minimizer has already seen a lower evaluation than a number the maximizing child sees.

Static evaluator

This function is used to evaluate a board position with regards to which player is winning, and by how much. The MiniMax algorithm will then choose the highest value for itself, while minimizing the options for its opponent. For gomoku, it was important to derive an evaluation function which could be calculated quickly, and which builds towards the final desired result of 5 squares in a row. Note that such a function would necessarily be isomorphic in four directions: vertical, horizontal, and on both diagonals.

My initial thought was that this would be extremely computationally expensive. There are many permutations of selected squares which can lead to a win, and many which do not. For example, XX — — OOO — — XX with O to move will lead to a win for O, but with X to move will not. However, I convinced myself that any static evaluation which built towards a win would find winning nodes at sufficient depth, so finding extremely detailed evaluation was less important than a general approximation.

Building from this thought, I decided to count the number of 4-in-a-rows (4s) and give them a high score, along with the 3s and 2s. Each in-a-row would be given an exponentially increasing “reward”, so that 4s scores much higher than 2×2s. For example, the payout function might be f(n) = 2^N for 2, 3, and 4 so that f(4) = 16 and 2×f(2) = 8. This ensures the desired result, that the optimal configuration of N squares is Ns.

Eventually, I determined that it was sufficient to simply count 2s with overlaps, since allowing double counts would still favor longer sequences of squares, but would not require separate checks for each length. Therefore, if 2s was rewarded 1, then XXX would be rewarded 2, and XXXX would be rewarded 3. This means that 4s is still more the most efficient configuration of four squares, since XX — XX only evaluates to 2.

Gameover function

This function simply needs to return true if the game is over and a player has won. After the simplifications to the static evaluator, the gameover function behaves almost identically. Instead of counting 2s, we check for the presence of 5s.

Bitmasks

Here is where the fun begins. I realized that a very efficient way of representing a game board was with a sequence of bits, where 1 represented an occupied square, and 0 represented an unoccupied square. A game state would therefore only require a bit sequence for each player (the game engine would prevent overlapping bits). For a 15×15=225 square board, each player’s occupied squares could be represented with a number 225 bits in length. Although Javascript Numbers are only 53 bits long, Javascript has a newer primitive, BigInt, which can store numbers of arbitrary length.

The biggest benefit of representing the game board this way is that it facilitates bitwise operations, which drastically reduces the time complexity for the static evaluator and gameover functions.

About BigInt

The BigInt primitive is a newer built-in type in Javascript, and as such is unsupported in some browsers. In particular, Internet Explorer and Safari do not have BigInt as a primitive. Although there are polyfills available for BigInt, they do not have the same performance as the native type. I decided that as a demonstration of the Minimax algorithm, supporting all browsers was not a priority.

Web Workers

Most people know that Javascript is single threaded. It is, except when it isn’t. Web Workers are a way of multithreading in the browser, which in this context is pretty important because it helps to avoid freezing the user interface. In this game, the board state is handed off to a Web Worker thread, which computes the best move and returns it to the main thread. Progress is reported back periodically to the main thread as well, which is shown in a progress bar underneath the Gomoku2049 logo.

Theoretically, I could have taken further advantage of multithreading when creating this game. Each branch in the decision tree can be parallelized, allowing for simultaneous computation of each node’s value. For example, a new thread could be used to evaluate each of the AI’s possible moves. Unfortunately, the number of possible moves for the AI can be quite high later in the game, and browsers limit the number of Web Workers allowed (Chrome allows 60, Firefox allows 20, etc.) so instead of spawning a new worker for each top level branch, threads would need to be spawned from a shared thread pool.

The full source code for this game can be found here.