LatentSeas: Exploring the Latent Space of Language Models

This blog post is fully human written. This blog post was also uploaded to substack to verify that it was human written

I've always been fascinated by the mechanistic interpretability work Anthropic does, and I found this paper especially interesting: Mapping the Mind of a Language Model. In short, models have patterns of activations that can be mapped to features, which get triggered when certain concepts or objects are present in the text. The article mentions, for example, a Golden Gate Bridge feature. Boosting the activations corresponding to that feature generates outputs that more frequently mention the boosted feature. I also found that these features have many parallels to Engram neurons in the brain (which was part of the inspiration for my work with the Machine Intelligence Lab at UofT).

Every concept and potential output lies dormant within the model's latent manifold. "Prompt Engineering," chain-of-thought, and agent orchestration are just workaround methods by which we can traverse this space. I wanted to explore directly traversing the internal activation space, and being able to directly manipulate it as well.

LatentSeas is that effort. It allows you to traverse the latent space of any open source model in three dimensions. It is a world of concepts which you can explore and manipulate.

How it works

Sparse Auto Encoders (SAEs)

An autoencoder is comprised of an encoder and a decoder. The encoder compresses some input A into a representation Z in a lower-dimensional space. The decoder then reconstructs A given only Z, and the goal is to minimize the reconstruction loss between the two.

Sparse Autoencoders instead upsample the input into a higher-dimensional space, then select the top-K features to build the sparse representation Z. The benefit is two-fold. First, the forced sparsity via the top-K selection guarantees that Z has k active features, and the increase in the number of overall features makes each individual feature more meaningful. Second, each feature is more likely to represent a unique concept, making it more interpretable and easier to understand. Typically, neurons in an LLM exhibit polysemanticity, meaning that multiple, unrelated concepts reside within the same neuron because the model must compress more concepts than it has dimensions for.

For the purposes of this project, I used Neuronpedia's SAE, which had directly annotated features. This SAE was trained on the activations of the residual stream of the transformer. These features are the core of the 3D world, where each feature is a point you can explore and probe. Moreover, you can manipulate the feature activations themselves to steer the model output.

World construction

With these annotated features from the SAE, I used UMAP to project the features into a 2D space. The height of any given location represents the feature density, which is calculated via a KDE (s/o STA238, I never thought I would use this knowledge again). Regions that are higher in altitude are denser in the number of related features. I then used K-means clustering to identify separate regions, to make the world more interesting to navigate.

Why Minecraft?

I initially wanted to make this in the browser, but it became more accessible and interesting as something I could explore in Minecraft instead. The goal was to build a navigable and manipulatable world, and I thought of Minecraft as the most interesting medium through which I could do it.

Results

Walkthrough

The latent space of GPT-2 in Minecraft.

You can probe the world to discover where certain features are: Typing the probe command. This lets me search for a specific concept and I can navigate toward where it exists in the latent space.

This takes us to the feature Result of navigating to the feature. Here, I've found the dog feature

I can then boost the activation of the feature that I found Boosting feature activations. This changes how strongly the dog feature fires.

Subsequent generations are now steered towards dogs The model now generates text that talks about dogs

However, since we're actually in the latent space of the model, we can find other related features within the vicinity of this dog feature. For example: In the area surrounding the dog feature, we have references to the concept of animals in general and zoos. We gain insight into how the model understands the relationship between these concepts and groups them together

I ran the boosted generation ten times, and all ten generations were related to dogs.

Interesting features

Large Empty Space

While exploring, I discovered a large swath of space devoid of features.

Huge gap in the latent space. No features around here at all, just flat, empty terrain for a long stretch.

There seems to be a little mound in this space. Uncovering the features there, I found:

Every feature here seems to be related to the word "the," in a little mound in the middle of nowhere.

It seems like there's a part of the model's brain solely dedicated to the word "the," separate from other places in the model, and there are multiple distinct features that all activate in response to the word "the." This is a known possible artifact of UMAP projections, which can distort global distances because it uses neighborhood structure for its mapping.

End of the World

At the end of the latent space, I found features corresponding to the end of text/documents. Coincidence? Maybe.

At the edge of the map, the features clustered here are about the end of documents and text.

What's Next

Currently, I'm only using GPT-2. The latent space of newer open source models (Gemma 2/3, Llama models) would be much cooler with more expressive feature sets. The generations from these models would also be more coherent, as GPT-2 is quite a small model. Unfortunately, I cannot run these larger models locally. I also want to make it more gamelike, potentially making it so every half-Minecraft-day the LLM spawns an object in your inventory based on how you steered the model, and you have to beat the game that way.