Minetester: A fully open RL environment built on Minetest

![](/images/blog/minetester-introduction/image10.png#center)

In the past several months we’ve seen great strides in the development of language models, especially in the private sector. Last year we saw several Minecraft-based environments released, including MineDojo and agents built on top of Minecraft such as VPT. This year we’ve seen agents that build on top of SOTA LLMs and libraries along with existing libraries, such as STEVE-1, Voyager, and GITM.

All of these are built on top of Minecraft and usually the MineRL environment. To meet our own research needs, we’ve been building a separate stack on top of the Minetest voxel engine.

What is Minetest, and why build an RL environment around it?[#](#what-is-minetest-and-why-build-an-rl-environment-around-it)
============================================================================================================================

The primary motivation for this project is to do and enable alignment research with AI agents that falls outside of the current language modelling paradigm. While progress for Minecraft-based RL agents has been significant, we find there are a few limitations to existing frameworks that make it more difficult for us to study the problems we’re interested in.

The first is a practical matter of transparency and software support. Existing frameworks are built on stacks involving multiple projects. For instance, with MineRL the stack looks as follows:

MineRL -> Malmo -> Minecraft -> Minecraft dependencies -> (???)

Minecraft itself is a closed-source game and must be decompiled for modding. Malmo is an old framework by Microsoft, and the original framework has not been updated for half a decade. This is not ideal from a software engineering and maintainability perspective. Adding new functionality requires reverse engineering and modifying quite a bit of old code.

This is in contrast to Minetest, where the entire stack is fully transparent, the codebase is smaller, and every dependency is open-source. This makes it easy for devs to dig into any component that could be causing an issue and lets us put our hooks directly into the Minetest game itself without unnecessary layers of indirection.

Another issue revolves around customizability: in the long term, we intend to support modes of operation in Minetest that are not really possible with existing frameworks. This includes synchronous and asynchronous multi-agent and human-agent interaction in game.

The last big reason is first-class modding support in Minetest. Minetest bills itself not as a game but more as a “Voxel Engine”, and as such, it is highly customizable. This makes it much easier to add modifications to the game physics to support specific experiments.

How does Minetester work under the hood?[#](#how-does-minetester-work-under-the-hood)
=====================================================================================

Minetester extends the standard Minetest environment in a few ways centred around enabling RL research.

![How the Minetester framework integrates with the Minetest game.](/images/blog/minetester-introduction/image5.png#center)
How the Minetester framework integrates with the Minetest game.

**Auxiliary data and communication (reward, feedback, etc.)**

The most important feature of Minetester is its extension of the modding API. This makes it possible to enable reward information to be relayed back to the client programmatically.

**Client-server synchronisation**

To keep the RL problem deterministic and the dynamics consistent, Minetester sends additional information between the client and the server to enable both to operate in lockstep.

**Headless operation**

Minetester supports 2 methods of headless operation. The first uses the Xvfb virtual framebuffer to encapsulate the Minetester process, the second, a WIP compile-time solution that replaces the standard rendering backend with an SDL2 headless backend.

**Python client wrapper**

Finally, the overall system is encapsulated in a Python wrapper that serves as an interface to AI agents built in modern ML frameworks.

Minetester baselines: PPO[#](#minetester-baselines-ppo)
=======================================================

To begin, we demonstrate basic usage of Minetester with a simple policy trained to just break wood. We started with PPO since it’s one of the simplest RL algorithms.

The first thing we noticed is that, without assistance, these algorithms don’t work at all, even for a “simple” task like breaking wood. By default, these simple algorithms will tend to flail around randomly. When they do manage to break a block, it’s very rare for it to be anything that would generate reward, so off-the-shelf algorithms don’t really work.

There are many different ways to deal with this. Ideally one would rely on more principled methods based on incentivising exploration and skill-building that naturally stumbles upon and quickly learns that breaking wood produces a reward. However, actually pulling this off is rather difficult, and we’re not aware of anyone successfully following this approach. Instead systems rely on other popular approaches that leverage prior knowledge and behavioural cloning. This was the approach taken by OpenAIs VPT and STEVE-1. Another tactic is just to make the problem easier. This is what DeepMind did with DreamerV3 and arguably what was done with GPT-4-based agents, which used APIs to dramatically simplify the action space.

In our case we opted to do something similar through a combination of reward-shaping and locking certain actions to further reduce the difficulty of the problem. Implementing both of these modifications in the Minetester framework is straightforward and simplifies the environment significantly.

On the agent side, we can simplify the action space and incentivise certain actions using standard modifiers to the gym environment. In our case we restrict the action space to just a few camera actions, moving left, right, and forward, and incentivising jumping with additional reward.

```

class AlwaysDig(MinetestWrapper):
 def step(self, action):
 action["DIG"] = True
 obs, rew, done, info = self.env.step(action)
 return obs, rew, done, info

```
*An example custom wrapper that makes the task easier by locking an action for the agent.*

On the environment side, we can implement more well-shaped reward functions by leveraging Minetest’s modding API. This lets us specify reward values for things like having a tree in frame, being in proximity to a tree, or breaking wood.

```
-- reward chopping tree nodes
minetest.register\_on\_dignode(function(pos, node)
 if string.find(node["name"], "tree") then
 REWARD = 1.0
 end
end)

```
Together these modifications make the problem tractable to learn for a simple PPO setup.

![A trained agent breaking wood.](https://cdn.discordapp.com/attachments/1072263460076929184/1125498773007777792/minetest_treechop.gif#center)
A trained agent breaking wood.

Interpreting the PPO baseline policy[#](#interpreting-the-ppo-baseline-policy)
==============================================================================

\**Note: The following section is best understood by following along with the notebook and model policy described [here](https://github.com/EleutherAI/minetest-interpretabilty-notebook).*

Even this simple policy contains interesting structure that we can deduce by inspecting the weights of the network and how it reacts to real data.

Whitebox interpretability without grounding[#](#whitebox-interpretability-without-grounding)
--------------------------------------------------------------------------------------------

We start by analysing the learned policy in a vacuum. This lets us make general statements about the structure of the network before we see how it reacts to real data.

**Activation probing/deep dreaming**
Since this is a visual model, we can copy some of the techniques from OAI’s [circuits publication](https://distill.pub/2020/circuits/zoom-in/). For instance, since the network we use is a ConvNet, we can use gradient ascent to probe what kind of image patches activate different neurons in the network. In our case, some quantities are particularly interpretable. These include actor-critic outputs and low-level neuron activations.

![Simplified NN diagram, 3 convolutional layers feed into value and critic heads. There are ReLU non-linearities between each layer. See the notebook for full implementation details.](/images/blog/minetester-introduction/image9.png#center)
Simplified NN diagram, 3 convolutional layers feed into value and critic heads. There are ReLU non-linearities between each layer. See the notebook for full implementation details.

This lets us ask questions like “What kind of images have high expected value?” and “What kind of images make the agent want to carry out a certain action, such as moving left/right?”

Doing this for high-value states is not super enlightening, but we do see a certain repeating low-level pattern show up.

![Image inputs with a high value according to the critic. We do see some repeating patterns but nothing very clear. Each image represents a different time delayed frame that gets fed into the NN.](/images/blog/minetester-introduction/image2.png#center)
Image inputs with a high value according to the critic. We do see some repeating patterns but nothing very clear. Each image represents a different time delayed frame that gets fed into the NN.

Another thing we can do is backprop through the probability that the agent preferentially turns to the right/left, which, going forward, we’ll call the yaw probability. These results are much less clean than what we see when deep dreaming with heavily trained classifiers. However, we can still make out some patterns. In particular, we see that the network is paying more attention to something closer to the middle-left of the screen when it wants to turn left, and stuff closer to the edges when it wants to turn right.

![“Saliency” of the most recent frame fed into the network. Left image represents turning left, the right image represents turning right. See the notebook for how these images were created from deep dream outputs.](/images/blog/minetester-introduction/image6.png#center)
“Saliency” of the most recent frame fed into the network. Left image represents turning left, the right image represents turning right. See the notebook for how these images were created from deep dream outputs.

This is very flimsy evidence, but it gives us a hypothesis for how the policy might work. The network always moves forward and jumps with some probability, but it needs to orient itself towards trees. Trees are “rough”, and if it sees “rough” on the right of its FOV, it orients itself to the right. Otherwise it does the opposite.

**Other observations**

We made a few other minor observations.

The first is that the network is clearly not symmetrical, which is surprising since the environment itself is mostly symmetrical. We suspect that this is mostly an artefact of training noise, but it’s possible that “symmetry breaking” in the policy is actually optimal.

The second is when you look at the matrices for the actor and the critic. The vectors are very much not random. Notably their dot-products are larger than random, indicating that both actor and critic heads are paying attention to the same features coming out of the base network.

Analysing real images and assessing actions values[#](#analysing-real-images-and-assessing-actions-values)
----------------------------------------------------------------------------------------------------------

Now that we have a hypothesis for what’s going on from looking at the network, we can see how the model reacts to real inputs to try to understand how the policy works in the environment. Since Minetest is a user-playable game, we can simply load it up and take some screenshots to feed into the network. One thing to note is that to make things faster and more computationally tractable, the input pipeline lowers the resolution of the inputs and strips away colour information. We can compare what the raw environment returns with what the network sees.

![What players see.](/images/blog/minetester-introduction/image3.png#center)
What players see.

![What the network sees.](/images/blog/minetester-introduction/image7.png#center)
What the network sees.

Due to downscaling and conversion to grayscale many details about the environment are lost to the network.

With that said, we can still take a look at how the network operates. We’re mainly interested in the yaw probability.

When feeding in real data we can very clearly confirm the network is implementing some kind of control system. This is made very clear by looking at how the yaw probability changes when we mirror an image or look at how the network reacts to a screenshot with a tree on the right or the left. This works with several different tree types and backgrounds.

![The yaw probability flips sign when we mirror the image.](/images/blog/minetester-introduction/image11.png#center)
The yaw probability flips sign when we mirror the image.

One thing we can check is how general this control system is. One way to do this is to evaluate the behaviour slightly out of distribution. Since the most straightforward hypothesis for how this network works is that it checks for brightness differences, we can either change textures or check how the network reacts to trees at night, where the contrast is inverted.

![This persists even at night, when trees are brighter than the environment.](/images/blog/minetester-introduction/image4.png#center)
This persists even at night, when trees are brighter than the environment.

The network still works. This rules out the possibility that the network is using something like light/dark on the right/left side of the screen to orient itself.

Drilling into the network[#](#drilling-into-the-network)
--------------------------------------------------------

The final and arguably hardest piece of the puzzle is to figure out how the network actually implements this control system in a way that explains our previous observations. This is rather difficult. To do so we drill into the network to better understand what its internals are doing. Since the network uses ReLU activations, we make the simplifying assumption that axis-aligned features are semantically meaningful. This is because ReLUs transform the data in an axis-aligned way. With this assumption, we can probe how convolutional neurons at each level of the network react to images.

**Layer 1**

The first layer contains simple linear filters, we can see a few different features, edge detectors, light detectors, and darkness detectors. This is straightforwardly visible by lining up the generated images with the regular images.

![An example of an edge detector in the first layer.](/images/blog/minetester-introduction/image1.png#center)
An example of an edge detector in the first layer.

**Layer 2**

With a single layer of non-linearity, the features tend to get more complicated. Overall the features tended to fall into 3 categories.

* The most common features were sensitive to the overall brightness of an image patch and would tend to either reflect or invert the brightness of the underlying image patch they observed.
* After this there was a large group of features that we were not able to effectively interpret. Their activation patterns did not clearly correspond to any feature in the images.
* Finally we did find 1 neuron that seemed to fairly consistently act as a tree detector.

![The “tree detector” neuron in the second layer.](/images/blog/minetester-introduction/image8.png#center)
The “tree detector” neuron in the second layer.

**Layer 3**

Neurons in layer 3 behaved approximately the same way, including a neuron that acted as a tree detector. To confirm that this feature played a role in the decisions made by the network, we computed the gradient w.r.t the yaw probability of the feature and found that activation of the neuron corresponded to a strong gradient in the direction that we expect.

Concrete takeaways[#](#concrete-takeaways)
------------------------------------------

Unfortunately, while we were able to identify individual components of the NN that corresponded to behaviours that we were interested in, we found diminishing returns to further drilling into the model, and we were not able to fully understand it. We can clearly see that the model actually does have reasonably fine-grained control of its actions. We can see it demonstrating advanced techniques, like not breaking the bottom log so that it can climb on top of it to reach and break more “branch” logs.

This investigation brought up several pain points in the workflow that we intend to improve going forward. Most of this revolves around our tooling.

The first is not having easy translation between what the user sees and does and what the network sees and does. The pipeline we’re using in the notebook was reconstructed by inspecting both the OAI Gym and the Minetester codebases, but ideally this would be done automatically.

The second is not having good facilities for recording user actions. For the purposes of this investigation, taking screenshots was sufficient to extract usable information, but as complexity ramps up, this will likely become insufficient.

The third is a general lack of tooling for interpretability research in the JAX ecosystem. For PyTorch, there are tools like TransformerLens, which makes interpretability of transformers easier, but as of this writing we’re unaware of any JAX equivalents.

More general/speculative takeaways[#](#more-generalspeculative-takeaways)
-------------------------------------------------------------------------

While these takeaways are less thoroughly supported by our investigation, they are high-level intuitions we’ve gathered from investigating the policy.

**There’s a big gap between a working NN and an easily interpretable one.**

Unlike with ConvNets for classification, this network seemed to have much less structure. Standard techniques for understanding them often showed much less structured information than what can be seen in something like VGG.

**In RL environments, understanding what the policy is doing is significantly easier with an understanding of the structure of the environment.**

The algorithm that the policy seems to implement is, at least on the surface, pretty simple: some circuitry to detect trees or associated features along with control circuitry that biases the agent to turn in one direction or the other, all the while encouraging the agent to move forward or jump. However, without at least a decent understanding of the game's mechanics, the effect of this policy is very difficult to determine. Even with a solid understanding of the game, without knowing the training objective ahead of time, and without access to entire episodes, determining the ultimate effects of the policy seems like it would be significantly more difficult.

It seems that perhaps model-based RL agents might be more interpretable, since they “internalise” their environment better. However, this is likely never going to work completely (since you can’t fit the whole universe inside your model), and other techniques will be necessary to understand how agents behave in environments we don’t fully understand ourselves.

**The structure of the network and the training algorithm plays a key role in facilitating interpretability.**

This seems to be a recurring theme with learned models. The actual underlying structure of the model and how it’s setup plays an important role in enabling interpretability. Things like induction heads are an emergent circuit in transformers due to the attention mechanism and the transformer architecture. Likewise, DeepDream-like visualisations in ConvNets are possible in part due to the restricted receptive fields and the continuous nature of their inputs. In our case, we exploited the convolutional structure and the relative simplicity and interpretability of our action/value mapping to at least partially reverse engineer the mechanics of the model.

Ultimately, it seems that interpretability techniques that work best for a given situation are sensitive to the architecture being studied.

**Interpretability and alignment should be first-class design considerations.**

Right now the way that we design and build AI systems is by iterating on the capabilities of those systems until they are essentially as powerful as we can make them, then going about trying to understand the inner workings of the model or how to align the model to behave the way we want.

Speculating a bit, it seems likely that this is a suboptimal way of going about things. If we intend to build systems that are interpretable and/or reliably aligned, alignment and interpretability will have to be first-class design considerations so that the necessary hooks to understand what’s going on are baked into the model.

Next Steps[#](#next-steps)
==========================

**Gymnasium: collaboration with the Farama Foundation and multi-agent support**

The current Minetester environment is built around the old OpenAI Gym API, but this is an outdated and unsupported API. The Farama Foundation has reached out to us, and going forward we plan to move to their more up-to-date Gymnasium API, which is a drop-in replacement with extra features and up-to-date maintenance. We also intend to work more closely with them on expanding the package further and adding new functionality, such as multi-agent support.

**Action recording and generative Minetest models**

While policy gradient is simple and straightforward to implement, it’s clearly limited. The next step we plan to take is to implement infrastructure for recording user actions. This will enable behavioural cloning and generative models of Minetest that are similar to DeepMind’s DreamerV3.

**Model based RL**

The long-term goal is to study embedded agency failure modes in the context of RL. As such, we plan to implement some MBRL baselines so that we can start studying how to interpret what they’re doing.

Join us![#](#join-us)
=====================

The Minetester project is large, and the number of different things to work on continues to grow. Checkout the #alignment-minetest project in our Discord to get involved. We have plenty of room for additional volunteers to contribute to different facets of the project.

Links[#](#links)
================

* [Minetester Repo](https://github.com/EleutherAI/minetest/)
* [Minetest Baselines](https://github.com/EleutherAI/minetest-baselines/)
* [Interpretabilty Notebook](https://github.com/EleutherAI/minetest-interpretabilty-notebook)

Curtis Huebner
Robert Klassert
Stepan Shabalin
Edwin Fennell
Delta Hessler


🐶Safetensors audited as really safe and becoming the default

Audit shows that safetensors is safe and ready to become the default
====================================================================

[HF中国镜像站](https://huggingface.co/), in close collaboration with [EleutherAI](https://www.eleuther.ai/) and [Stability AI](https://stability.ai/), has ordered
an external security audit of the `safetensors` library, the results of which allow
all three organizations to move toward making the library the default format
for saved models.

The full results of the security audit, performed by [Trail of Bits](https://www.trailofbits.com/),
can be found here: [Report](https://huggingface.co/datasets/safetensors/trail_of_bits_audit_repot/resolve/main/SOW-TrailofBits-EleutherAI_HuggingFace-v1.2.pdf).

The following blog post explains the origins of the library, why these audit results are important,
and the next steps.

What is safetensors?[#](#what-is-safetensors)
=============================================

🐶[Safetensors](https://github.com/huggingface/safetensors) is a library
for saving and loading tensors in the most common frameworks (including PyTorch, TensorFlow, JAX, PaddlePaddle, and NumPy).

For a more concrete explanation, we'll use PyTorch.

```
import torch
from safetensors.torch import load\_file, save\_file

weights = {"embeddings": torch.zeros((10, 100))}
save\_file(weights, "model.safetensors")
weights2 = load\_file("model.safetensors")

```
It also has a number of [cool features](https://github.com/huggingface/safetensors#yet-another-format-) compared to other formats, most notably that loading files is *safe*, as we'll see later.

When you're using `transformers`, if `safetensors` is installed, then those files will already
be used preferentially in order to prevent issues, which means that

```
pip install safetensors

```
is likely to be the only thing needed to run `safetensors` files safely.

Going forward and thanks to the validation of the library, `safetensors` will now be installed in `transformers` by
default. The next step is saving models in `safetensors` by default.

We are thrilled to see that the `safetensors` library is already seeing use in the ML ecosystem, including:

* [Civitai](https://civitai.com/)
* [Stable Diffusion Web UI](https://github.com/AUTOMATIC1111/stable-diffusion-webui)
* [dfdx](https://github.com/coreylowman/dfdx)
* [LLaMA.cpp](https://github.com/ggerganov/llama.cpp/blob/e6a46b0ed1884c77267dc70693183e3b7164e0e0/convert.py#L537)

Why create something new?[#](#why-create-something-new)
=======================================================

The creation of this library was driven by the fact that PyTorch uses `pickle` under
the hood, which is inherently unsafe. (Sources: [1](https://huggingface.co/docs/hub/security-pickle), [2, video](https://www.youtube.com/watch?v=2ethDz9KnLk), [3](https://github.com/pytorch/pytorch/issues/52596))

With pickle, it is possible to write a malicious file posing as a model
that gives full control of a user's computer to an attacker without the user's knowledge,
allowing the attacker to steal all their bitcoins 😓.

While this vulnerability in pickle is widely known in the computer security world (and is acknowledged in the PyTorch [docs](https://pytorch.org/docs/stable/generated/torch.load.html)), it’s not common knowledge in the broader ML community.

Since the HF中国镜像站 Hub is a platform where anyone can upload and share models, it is important to make efforts
to prevent users from getting infected by malware.

We are also taking steps to make sure the existing PyTorch files are not malicious, but the best we can do is flag suspicious-looking files.

Of course, there are other file formats out there, but
none seemed to meet the full set of [ideal requirements](https://github.com/huggingface/safetensors#yet-another-format-) our team identified.

In addition to being safe, `safetensors` allows lazy loading and generally faster loads (around 100x faster on CPU).

Lazy loading means loading only part of a tensor in an efficient manner.
This particular feature enables arbitrary sharding with efficient inference libraries, such as [text-generation-inference](https://github.com/huggingface/text-generation-inference), to load LLMs (such as LLaMA, StarCoder, etc.) on various types of hardware
with maximum efficiency.

Because it loads so fast and is framework agnostic, we can even use the format
to load models from the same file in PyTorch or TensorFlow.

The security audit[#](#the-security-audit)
==========================================

Since `safetensors` main asset is providing safety guarantees, we wanted to make sure
it actually delivered. That's why HF中国镜像站, EleutherAI, and Stability AI teamed up to get an external
security audit to confirm it.

Important findings:

* No critical security flaw leading to arbitrary code execution was found.
* Some imprecisions in the spec format were detected and fixed.
* Some missing validation allowed [polyglot files](https://en.wikipedia.org/wiki/Polyglot_(computing)), which was fixed.
* Lots of improvements to the test suite were proposed and implemented.

In the name of openness and transparency, all companies agreed to make the report
fully public.

[Full report](https://huggingface.co/datasets/safetensors/trail_of_bits_audit_repot/resolve/main/SOW-TrailofBits-EleutherAI_HuggingFace-v1.2.pdf)

One import thing to note is that the library is written in Rust. This adds
an extra layer of [security](https://doc.rust-lang.org/rustc/exploit-mitigations.html)
coming directly from the language itself.

While it is impossible to
prove the absence of flaws, this is a major step in giving reassurance that `safetensors`
is indeed safe to use.

Going forward[#](#going-forward)
================================

For HF中国镜像站, EleutherAI, and Stability AI, the master plan is to shift to using this format by default.

EleutherAI has added support for evaluating models stored as `safetensors` in their LM Evaluation Harness and is working on supporting the format in their GPT-NeoX distributed training library.

Within the `transformers` library we are doing the following:

* Create `safetensors`.
* Verify it works and can deliver on all promises (lazy load for LLMs, single file for all frameworks, faster loads).
* Verify it's safe. (This is today's announcement.)
* Make `safetensors` a core dependency. (This is already done or soon to come.)
* Make `safetensors` the default saving format. This will happen in a few months when we have enough feedback
to make sure it will cause as little disruption as possible and enough users already have the library
to be able to load new models even on relatively old `transformers` versions.

As for `safetensors` itself, we're looking into adding more advanced features for LLM training,
which has its own set of issues with current formats.

Finally, we plan to release a `1.0` in the near future, with the large user base of `transformers` providing the final testing step.
The format and the lib have had very few modifications since their inception,
which is a good sign of stability.

We're glad we can bring ML one step closer to being safe and efficient for all!

Nicolas Patry
Stella Biderman
Garry Jean-Baptiste


Alignment Research @ EleutherAI

The past and future of AI alignment at Eleuther[#](#the-past-and-future-of-ai-alignment-at-eleuther)
----------------------------------------------------------------------------------------------------

Initially, EleutherAI focused mainly on supporting open source research. AI alignment was something that was acknowledged by many of the core members as important, but it was not the primary focus. We mainly had discussions about the topic in the #alignment channel and other parts of our discord while we worked on other projects.

As EAI grew, AI alignment started to get taken more seriously, especially by its core members. What started off as a single channel turned into a whole host of channels about different facets of alignment. We also hosted several reading groups related to alignment, such as the modified version of [Richard Ngo’s curriculum](https://www.alignmentforum.org/posts/Zmwkz2BMvuFFR8bi3/agi-safety-fundamentals-curriculum-and-application) and an interpretability reading group. Eventually alignment became the central focus for a large segment of EAIs leadership, so much so that all our previous founders went off to do full time alignment research at [Conjecture](https://conjecture.dev/) and OpenAI.

Right now, the current leadership believes making progress in AI alignment is very important. The organization as a whole is involved in a mix of alignment research, interpretability work, and other projects that we find interesting.

Moving forward, EAI remains committed to facilitating and enabling open source research, and plans to ramp up its alignment and interpretability research efforts. We want to increase our understanding and control of modern ML systems and minimize existential risks posed by artificial intelligence.

Our meta-level approach to alignment[#](#our-meta-level-approach-to-alignment)
------------------------------------------------------------------------------

It is our impression that AI alignment is still a very pre-paradigmatic field. Progress in the field often matches the research pattern we see in the [ELK report](https://docs.google.com/document/d/1WwsnJQstPq91_Yh-Ch2XRL8H_EpsnjrC1dwZXR37PC8/edit), where high level strategies are proposed, and problems or counterexamples are found. Sometimes these issues can be fixed, but oftentimes fundamental issues are identified that make an initially promising approach less interesting.

A consequence of this is that it’s difficult to commit to an object level strategy to make progress on AI alignment, and even harder to commit to any grand strategic plan to solve the problem. Instead it makes more sense to have a meta level strategy that makes us better able to leverage our unique position within the AI research ecosystem, and pivot when we get new information.

Going forward, this means we want to pursue interesting projects that meet a few general desiderata.

1. Our volunteers, partners, and collaborators are enthusiastic about the project.
2. We believe that pursuing the project won’t lead to a net increase in existential risks from AI. We’ll check for this even if the project is ostensibly a project that will greatly increase our understanding of AI alignment.
3. The project is something that EAI is better equipped to do than anyone else in the space, or the project seems interesting or important, but neglected by the broader community.

In order to pull this off, we aim to stay on top of both the latest developments in AI and alignment research. We’ll also carefully consider new projects before we embark on or allocate resources for them.

Problems we are interested in and research we are doing right now[#](#problems-we-are-interested-in-and-research-we-are-doing-right-now)
----------------------------------------------------------------------------------------------------------------------------------------

Given the current state AI landscape, there are a few directions that we find especially interesting. We’d love to collaborate with others to make progress on these issues.

#### Interpretability work[#](#interpretability-work)

Interpretability work, especially with current models, seems like a very tractable and scalable research direction. It seems especially easy for current ML researchers to pick up and make progress on it. EAI is well equipped to enable this kind of research, especially for larger language models that more closely resemble the ones we see in modern production systems.

This is arguably where most of our recent efforts have been lately, as exemplified by projects like the #interpreting-across-time and #interpreting-across-depth projects and their associated channels.

#### Value specification in embedded agents[#](#value-specification-in-embedded-agents)

This is another core problem that we think is very important, and we mention it as a key difficulty in our second retrospective. In practice any mechanism for aligning AI systems, or correcting their behavior after they’ve been deployed, is going to be part of the same environment as the agent, and prone to it’s interference. This shows up classically in the wireheading scenario, where an AI hijacks it’s own reward signal instead of doing whatever it was that we originally wanted, but there’s reason to believe that similar problems might show up with more sophisticated value specification schemes. While we haven’t seen any issues in deployed systems related to this issue, it’s something we’re worried might show up down the line with more competent/powerful systems.

Tackling this problem and other related issues is the long term goal of our alignment-minetest project, which aims to create a sandbox that we can use to study the embedded system failures, amongst other issues.

#### Other directions related to alignment that we find promising[#](#other-directions-related-to-alignment-that-we-find-promising)

There are a bunch of other directions and problems we find worth studying. An incomplete list includes:

* **Eliciting latent knowledge:** The [ELK report](https://docs.google.com/document/d/1WwsnJQstPq91_Yh-Ch2XRL8H_EpsnjrC1dwZXR37PC8/edit) highlights a problem that seems to at least partially capture a core difficulty in alignment research.
* **Practical implementation of alignment schemes:** There have been a whole host of alignment schemes proposed over the years, and much theoretical debate about the merits of each approach. We think it’s important to test implementations of these schemes even if we know they won't work in the limit. This is to validate them and uncover problems that only show up in practice.
* **A better theoretical understanding of model misspecification:** Bayesian inference is a powerful lens that can be used to understand modern AI systems, and it gives us a way to understand “ideal” AI systems that are given access to infinite amounts of computing power. However, real systems are not ideal. One of the core assumptions of Bayesian inference, realizability, is violated in pretty much every real system we deploy. We want to understand how these systems “just work anyways”, and to the extent that they don’t, why and how to fix them.
* **Verifiable computation:** better protocols for verifying and securing computation at the hardware, software, and conceptual levels would be great steps forward for providing theoretical bounds on system capabilities.


Transformer Math 101

Introduction[#](#introduction)
==============================

A lot of basic, important information about transformer language models can be computed quite simply. Unfortunately, the equations for this are not widely known in the NLP community. The purpose of this document is to collect these equations along with related knowledge about where they come from and why they matter.

**Note:** This post is primarily concerned with training costs, which are dominated by VRAM considerations. For an analogous discussion of inference costs with a focus on latency, check out [this excellent blog post](https://kipp.ly/blog/transformer-inference-arithmetic/) by Kipply.

Compute Requirements[#](#compute-requirements)
==============================================

The basic equation giving the cost to train a transformer model is given by:

$$
C\approx\tau T = 6PD
$$

where:

* $C$ is the compute required to train the transformer model, in total floating point operations
* $C=C\_{\text{forward}}+C\_{\text{backward}}$
* $C\_{\text{forward}}\approx2PD$
* $C\_{\text{backward}}\approx4PD$
* $\tau$ is the aggregate throughput of your hardware setup ($\tau=(\text{No. GPUs}) \times (\text{Actual FLOPs}/\text{GPU})$), in FLOPs
* $T$ is the time spent training the model, in seconds
* $P$ is the number of parameters in the transformer model
* $D$ is the dataset size, in tokens

These equations are proposed and experimentally validated in [OpenAI’s scaling laws paper](https://arxiv.org/abs/2001.08361) and [DeepMind’s scaling laws paper](https://arxiv.org/abs/2203.15556). Please see each paper for more information.

It’s worth taking an aside and discussing the units of $C$. $C$ is a measure of total compute, but can be measured by many units such as:

* FLOP-seconds, which is in units of $[\frac{\text{Floating Point Operations}}{\text{Second}}] \times [\text{Seconds}]$
* GPU-hours, which is in units of $[\text{No. GPUs}]\times[\text{Hours}]$
* Scaling laws papers tend to report values in PetaFLOP-days, or $10^{15}\times24\times3600$ total floating point operations

One useful distinction to keep in mind is the concept of $\text{Actual FLOPs}$. While GPU accelerator whitepapers usually advertise their theoretical FLOPs, these are never met in practice (especially in a distributed setting!). Some common reported values of $\text{Actual FLOPs}$ in a distributed training setting are reported below in the Computing Costs section.

Note that we use the throughput-time version of the cost equation as used in [this wonderful blog post on LLM training costs](https://medium.com/@dzmitrybahdanau/the-flops-calculus-of-language-model-training-3b19c1f025e4).

Parameter vs Dataset Tradeoffs[#](#parameter-vs-dataset-tradeoffs)
------------------------------------------------------------------

Although strictly speaking you can train a transformer for as many tokens as you like, the number of tokens trained can highly impact both the computing costs and the final model performance making striking the right balance important.

**Let’s start with the elephant in the room: “compute optimal” language models.** Often referred to as “Chinchilla scaling laws” after the model series in the paper that gave rise to current beliefs about the number of parameters, a compute optimal language model has a **number of parameters** and a **dataset size** that satisfies the approximation $D=20P$. This is optimal in one very specific sense: in a resource regime where using 1,000 GPUs for 1 hour and 1 GPU for 1,000 hours cost you the same amount, if your goal is to maximize performance while minimizing the cost in GPU-hours to train a model you should use the above equation.

**We do not recommend training a LLM for less than 200B tokens.** Although this is “chinchilla optimal” for many models, the resulting models are typically quite poor. For almost all applications, we recommend determining what inference cost is acceptable for your usecase and training the largest model you can to stay under that inference cost for as many tokens as you can.

Engineering Takeaways for Compute Costs[#](#engineering-takeaways-for-compute-costs)
------------------------------------------------------------------------------------

Computing costs for transformers are typically listed in GPU-hours or FLOP-seconds.

* GPT-NeoX achieves 150 TFLOP/s/A100 with normal attention and 180 TFLOP/s/A100 with Flash Attention. This is in line with other highly optimized libraries at scale, for example Megatron-DS reports between 137 and 163 TFLOP/s/A100.
* As a general rule of thumb, you should always be able to achieve approximately 120 TFLOP/s/A100. If you are seeing below 115 TFLOP/s/A100 there is probably something wrong with your model or hardware configuration.
* With high-quality interconnect such as InfiniBand, you can achieve linear or sublinear scaling across the data parallel dimension (i.e. increasing the data parallel degree should increase the overall throughput nearly linearly). Shown below is a plot from testing the GPT-NeoX library on Oak Ridge National Lab’s Summit supercomputer. Note that V100s are on the x-axis, while most of the numerical examples in the post are for A100s.

![GPT-NeoX scaling](/images/blog/transformer-math/neox-scaling.png#center)

Memory Requirements[#](#memory-requirements)
============================================

Transformers are typically described in terms of their *size in parameters*. However, when determining what models can fit on a given set of computing resources you need to know **how much space in bytes** the model will take up. This can tell you how large a model will fit on your local GPU for inference, or how large a model you can train across your cluster with a certain amount of total accelerator memory.

Inference[#](#inference)
------------------------

### Model Weights[#](#model-weights)

![Model Weights](https://cdn.discordapp.com/attachments/938462108721483787/1052372619577532467/image.png#center)

Most transformers are trained in **mixed precision**, either fp16 + fp32 or bf16 + fp32. This cuts down on the amount of memory required to train the models, and also the amount of memory required to run inference. We can cast language models from fp32 to fp16 or even int8 without suffering a substantial performance hit. These numbers refer to the size *in bits* a single parameter requires. Since there are 8 bits in a Byte, we divide this number by 8 to see how many Bytes each parameter requires

* In int8, $\text{memory}\_{\text{model}}=(1 \text{ byte} /\text{param})\cdot ( \text{No. params})$
* In fp16 and bf16, $\text{memory}\_{\text{model}}=(2 \text{ bytes} /\text{param})\cdot ( \text{No. params})$
* In fp32, $\text{memory}\_{\text{model}}=(4 \text{ bytes} /\text{param})\cdot (\text{No. params})$

There is also a small amount of additional overhead, which is typically irrelevant to determining the largest model that will fit on your GPU. In our experience this overhead is ≤ 20%.

### Total Inference Memory[#](#total-inference-memory)

In addition to the memory needed to store the model weights, there is also a small amount of additional overhead during the actual forward pass. In our experience this overhead is ≤ 20% and is typically irrelevant to determining the largest model that will fit on your GPU.

In total, a good heuristic answer for “will this model fit for inference” is:

$\text{Total Memory}\_{\text{Inference}}\approx(1.2) \times \text{Model Memory}$

We will not investigate the sources of this overhead in this blog post and leave it to other posts or locations for now, instead focusing on memory for model training in the rest of this post. If you’re interested in learning more about the calculations required for inference, check out [this fantastic blog post covering inference in depth](https://kipp.ly/blog/transformer-inference-arithmetic/). Now, on to training!

Training[#](#training)
----------------------

In addition to the model parameters, training requires the storage of optimizer states and gradients in device memory. This is why asking “how much memory do I need to fit model X” immediately leads to the answer “this depends on training or inference.” Training always requires more memory than inference, often very much more!

### Model Parameters[#](#model-parameters)

First off, models can be trained in pure fp32 or fp16:

* Pure fp32, $\text{memory}\_{\text{model}}=(4 \text{ bytes} /\text{param})\cdot (\text{No. params})$
* Pure fp16, $\text{memory}\_{\text{model}}=(2 \text{ bytes} /\text{param})\cdot (\text{No. params})$

In addition to the common model weight datatypes discussed in Inference, training introduces **mixed-precision** training such as [AMP](https://developer.nvidia.com/automatic-mixed-precision). This technique seeks to maximize the throughput of GPU tensor cores while maintaining convergence. The modern DL training landscape frequently uses mixed-precision training because: 1) fp32 training is stable, but has a high memory overhead and doesn’t exploit NVIDIA GPU tensor cores, and 2) fp16 training is stable but difficult to converge. For more information on mixed-precision training, we recommend reading [this notebook by tunib-ai](https://nbviewer.org/github/tunib-ai/large-scale-lm-tutorials/blob/main/notebooks/08_zero_redundancy_optimization.ipynb). Note that mixed-precision requires an fp16/bf16 and fp32 version of the model to be stored in memory, requiring:

* Mixed-precision (fp16/bf16 and fp32), $\text{memory}\_{\text{model}}=(2 \text{ bytes} /\text{param})\cdot (\text{No. params})$

plus an additional size $(4\text{ bytes/param}) \cdot (\text{#params})$ copy of the model **that we count within our optimizer states**.

### Optimizer States[#](#optimizer-states)

Adam is magic, but it’s highly memory inefficient. In addition to requiring you to have a copy of the model parameters and the gradient parameters, you also need to keep an additional three copies of the gradient parameters. Therefore,

* For vanilla AdamW, $\text{memory}\_{\text{optimizer}}=(12 \text{ bytes}/\text{param})\cdot (\text{No. params})$
	+ fp32 copy of parameters: 4 bytes/param
	+ Momentum: 4 bytes/param
	+ Variance: 4 bytes/param
* For 8-bit optimizers like [bitsandbytes](https://github.com/TimDettmers/bitsandbytes), $\text{memory}\_{\text{optimizer}}=(6 \text{ bytes} /\text{param})\cdot (\text{No. params})$
	+ fp32 copy of parameters: 4 bytes/param
	+ Momentum: 1 byte/param
	+ Variance: 1 byte/param
* For SGD-like optimizers with momentum, $\text{memory}\_{\text{optimizer}}=(8 \text{ bytes} /\text{param})\cdot (\text{No. params})$
	+ fp32 copy of parameters: 4 bytes/param
	+ Momentum: 4 bytes/param

### Gradients[#](#gradients)

Gradients can be stored in fp32 or fp16 (Note that the gradient datatype often matches the model datatype. We see that it therefore is stored in fp16 for fp16 mixed-precision training), so their contribution to memory overhead is given by:

* In fp32, $\text{memory}\_{\text{gradients}}=(4 \text{ bytes} /\text{param})\cdot (\text{No. params})$
* In fp16, $\text{memory}\_{\text{gradients}}=(2 \text{ bytes} /\text{param})\cdot (\text{No. params})$

### Activations and Batch Size[#](#activations-and-batch-size)

Modern GPUs are typically bottlenecked by memory, not FLOPs, for LLM training. Therefore activation recomputation/checkpointing is an extremely popular method of trading reduced memory costs for extra compute costs. Activation recomputation/checkpointing works by recomputing activations of certain layers instead of storing them in GPU memory. The reduction in memory depends on how selective we are when deciding which layers to clear, but Megatron’s selective recomputation scheme is depicted in the figure below:

![activation memory](/images/blog/transformer-math/activations.png#center)

Where the dashed red line indicates the memory capacity of an A100-80GB GPU, and “present work” indicates the memory requirements after applying selective activation recomputation. See [Reducing Activation Recomputation in Large Transformer Models](https://arxiv.org/abs/2205.05198) for further details and the derivation of the equations below

The basic equation giving the memory required to store activations for a transformer model is given by:

$$
\begin{align\*}\text{memory}^{\text{No Recomputation}}\_{\text{activations}}=sbhL(10+\frac{24}{t}+5\frac{a \cdot s}{h\cdot t}) \text{ bytes}\end{align\*}
$$

$$
\begin{align\*}\text{memory}^{\text{Selective Recomputation}}\_{\text{activations}}=sbhL(10+\frac{24}{t}) \text{ bytes}\end{align\*}
$$

$$
\begin{align\*}\text{memory}^{\text{Full Recomputation}}\_{\text{activations}}=2 \cdot sbhL \text{ bytes}\end{align\*}
$$

where:

* $s$ is the sequence length, in tokens
* $b$ is the batch size per GPU
* $h$ is the dimension of the hidden size within each transformer layer
* $L$ is the number of layers in the transformer model
* $a$ is the number of attention heads in the transformer model
* $t$ is the degree of tensor parallelism being used (1 if not)
* We assume no sequence parallelism is being used
* We assume that activations are stored in fp16

The additional recomputation necessary also depends on the selectivity of the method, but it’s bounded above by a full additional forward pass. Hence the updated cost of the forward pass is given by:

$$
2PD\leq C\_{\text{forward}}\leq4PD
$$

### Total Training Memory[#](#total-training-memory)

Therefore, a good heuristic answer for “will this model fit for training” is:

$$
\begin{align\*}\text{Total Memory}\_{\text{Training}} = \text{Model Memory}+\text{Optimiser Memory}+\text{Activation Memory}+\text{Gradient Memory}\end{align\*}
$$

Distributed Training[#](#distributed-training)
----------------------------------------------

### Sharded Optimizers[#](#sharded-optimizers)

The massive memory overheads for optimizers is the primary motivation for sharded optimizers such as [ZeRO](https://arxiv.org/abs/1910.02054) and [FSDP](https://engineering.fb.com/2021/07/15/open-source/fsdp/). Such sharding strategies reduce the optimizer overhead by a factor of $\text{No. GPUs}$, which is why a given model configuration may fit at large scale but OOM at small scales. If you’re looking to calculate the memory overhead required by training using a sharded optimizer, you will need to include the equations from the figure below. For some sample calculations of sharded optimization, see the following figure from the [ZeRO](https://arxiv.org/abs/1910.02054) paper (Note that $P\_{os}$ $P\_{os+g}$ and $P\_{os+g+p}$ are commonly denoted as ZeRO-1, ZeRO-2, ZeRO-3, respectively. ZeRO-0 commonly means “ZeRO disabled”):

![ZeRO illustration](/images/blog/transformer-math/zero_fig.png#center)

![ZeRO legend](/images/blog/transformer-math/zero_legend.png#center)

In the language of this blog post (assuming mixed-precision and the Adam optimizer):

* For ZeRO-1,

$$
\begin{align\*}\text{Total Memory}\_{\text{Training}}\approx\text{Model Memory}+\frac{\text{Optimizer memory}}{(\text{No. GPUs})}+\text{Activation Memory}+\text{Gradient Memory}\end{align\*}
$$

* For ZeRO-2,

$$
\begin{align\*}\text{Total Memory}\_{\text{Training}}\approx\text{Model Memory}+\text{Activation Memory}+\frac{\text{Optimizer Memory}+\text{Gradient Memory}}{(\text{No. GPUs})}\end{align\*}
$$

* For ZeRO-3,

$$
\begin{align\*}\text{Total Memory}\_{\text{Training}}\approx \text{Activation Memory}+\frac{\text{Model Memory}+\text{Optimizer Memory}+\text{Gradient Memory}}{(\text{No. GPUs})} + \text{(ZeRO-3 Live Params)}\end{align\*}
$$

Where $(\text{DP Degree})$ is just $(\text{No. GPUs})$ unless pipeline and/or tensor parallelism are applied. See [Sharded Optimizers + 3D Parallelism](https://www.notion.so/Sharded-Optimizers-3D-Parallelism-9c476d020d7641a299fb6be6ae82e9f8) for details.

Note that ZeRO-3 introduces a set of live parameters. This is because ZeRO-3 introduces a set of config options (***stage3\_max\_live\_parameters, stage3\_max\_reuse\_distance, stage3\_prefetch\_bucket\_size, stage3\_param\_persistence\_threshold***) that control how many parameters are within GPU memory at a time (larger values take more memory but require less communication). Such parameters can have a significant effect on total GPU memory.

Note that ZeRO can also partition activations over data parallel ranks via **ZeRO-R**. This would also bring the $\text{memory}\_\text{activations}$ above the tensor parallelism degree $t$. For more details, read the associated [ZeRO paper](https://arxiv.org/abs/1910.02054) and [config options](https://www.deepspeed.ai/docs/config-json/#activation-checkpointing) (note in GPT-NeoX, this is the `partition_activations` flag). If you are training a huge model, you would like to trade some memory overhead for additional communication cost, and activations become a bottleneck. As an example of using ZeRO-R along with ZeRO-1:

$$
\begin{align\*}\text{Total Memory}\_{\text{Training}}\approx\text{Model Memory}+\frac{\text{Optimizer Memory}}{(\text{No. GPUs})}+\frac{\text{Activation Memory}}{\text{(Tensor-Parallel-Size)}}+\text{Gradient Memory}\end{align\*}
$$

### 3D Parallelism[#](#3d-parallelism)

Parallelism for LLMs comes in 3 primary forms:

**Data parallelism:** Split the data among (possibly model-parallel) replicas of the model

**Pipeline or Tensor/Model parallelism:** These parallelism schemes split the parameters of the model across GPUs. Such schemes require significant communication overhead, but their memory reduction is approximately:

$$
\begin{align\*}\text{memory}^{\text{w/ parallelism}}\_{\text{model}}\approx\frac{\text{Model Memory}}{\text{(Pipe-Parallel-Size})\times\text{(Tensor-Parallel-Size)}}\end{align\*}
$$

$$
\begin{align\*}\text{memory}^{\text{w/ parallelism}}\_{\text{gradients}}\approx\frac{\text{Gradient Memory}}{\text{(Pipe-Parallel-Size})}\end{align\*}
$$

Note that this equation is approximate due to the facts that (1) pipeline parallelism doesn’t reduce the memory footprint of activations, (2) pipeline parallelism requires that all GPUs store the activations for all micro-batches in-flight, which becomes significant for large models, and (3) GPUs need to temporarily store the additional communication buffers required by parallelism schemes.

### Sharded Optimizers + 3D Parallelism[#](#sharded-optimizers--3d-parallelism)

When ZeRO is combined with tensor and/or pipeline parallelism, the resulting parallelism strategy forms a mesh like the following:

![3D parallelism](https://i.imgur.com/xMgptTN.png#center)

As an important aside, the DP degree is vital for use in calculating the global batch size of training. The data-parallel degree depends on the number of complete model replicas:

$$
\begin{align\*}\text{DP Degree = }\frac{\text{No. GPUs}}{\text{(Pipe-Parallel-Size})\times\text{(Tensor-Parallel-Size)}}\end{align\*}
$$

While pipeline parallelism and tensor parallelism are compatible with all stages of ZeRO (e.g. ZeRO-3 with tensor parallelism would lead to us first slicing the tensors, then applying ZeRO-3 within each tensor-parallel unit), only ZeRO-1 tends to perform well in conjunction with tensor and/or pipeline parallelism. This is due to the conflicting parallelism strategies for gradients (pipeline parallelism and ZeRO-2 both split gradients) (tensor parallelism and ZeRO-3 both split model parameters), which leads to a significant communication overhead.

Putting everything together for a typical 3D-parallel ZeRO-1 run with activation partitioning:

$$
\begin{align\*}\text{Total Memory}\_{\text{Training}}\approx\frac{\text{Model Memory}}{\text{(Pipe-Parallel-Size})\times\text{(Tensor-Parallel-Size)}}+\frac{\text{Optimizer Memory}}{(\text{No. GPUs})}+\frac{\text{Activation Memory}}{\text{(Tensor-Parallel-Size)}}+\frac{\text{Gradient Memory}}{\text{(Pipe-Parallel-Size})}\end{align\*}
$$

Conclusion[#](#conclusion)
==========================

EleutherAI engineers frequently use heuristics like the above to plan efficient model training and to debug distributed runs. We hope to provide some clarity on these often-overlooked implementation details, and would love to hear your feedback at [contact@eleuther.ai](mailto:contact@eleuther.ai) if you would like to discuss or think we’ve missed anything!

To cite this blog post, please use:

```
@misc{transformer-math-eleutherai,
 title = {Transformer Math 101},
 author = {Anthony, Quentin and Biderman, Stella and Schoelkopf, Hailey},
 howpublished = \url{blog.eleuther.ai/},
 year = {2023}
}

```

Quentin Anthony
Stella Biderman
Hailey Schoelkopf


Exploratory Analysis of TRLX RLHF Transformers with TransformerLens

Introduction[#](#introduction)
------------------------------

LLMs trained with RLHF are a prominent paradigm in the current AI landscape, yet not much mechanistic interpretability work has been done on these models to date--partially due to the complexity and scale of these models, and partially due to the previous lack of accessible tooling for training and analysis.

Fortunately, we are reaching the point where tooling for both mechanistic interpretability and for RLHF fine-tuning is becoming available. In this blog post, I demonstrate how to do both RLHF training using TRLX, an open-source library created by CarperAI; and mechanistic interpretation of TRLX models using TransformerLens, a library created by Neel Nanda. Rather than going deep into specific findings, I want to illustrate some processes and tools I think are useful.

**This post is intended to summarize and go alongside an interactive Colab;**[you can find that here](https://colab.research.google.com/drive/1DK6_HNRjUHliolQ2uMYNpB24XyeD9BIl).

I first fine-tune a movie-review-generating version of GPT-2 with TRLX to generate only negatively-biased movie reviews, following an example provided in the TRLX repo. I then load and analyze the model (and the original model before RLHF) into TransformerLens for mechanistic interpretability analysis. Here, I adapt some of the techniques and code from Neel Nanda's excellent [Exploratory Analysis Demo](https://colab.research.google.com/github/neelnanda-io/Easy-Transformer/blob/main/Exploratory_Analysis_Demo.ipynb).

In addition to carrying out some basic analysis to understand how different layers contribute to the logits, I also identify some key regions of the network responsible for contributing the negative bias to the network (at least, for the specific task of predicting the next adjective). Much analysis remains to be done, but I hope this work provides a useful starting point.

### Importance of RLHF[#](#importance-of-rlhf)

RLHF (or sometimes, RLAIF, or RL from AI Feedback) is becoming increasingly important as a method for specifying the behavior of LLMs like OpenAI's ChatGPT or Anthropic's Claude. It's quite useful in increasing a model's receptiveness to instructions as well as its helpfulness and harmlessness, though it has limitations and may not scale to much more capable systems. Nevertheless, it is quite important in today's LLM landscape.

RL induces behavior in models that are critical to understand as we delegate more tasks to them. Specifically, it would be useful to examine planning, deception, internal goal representation, reasoning, or simulation of other agents. Neel Nanda provides a set of [recommended RL problems](https://www.lesswrong.com/s/yivyHaCAmMJ3CqSyj/p/eqvvDM25MXLGqumnf) in his 200 Open Problems in Mechanistic Interpretability sequence. In this notebook, the process I outline (of breaking things down to small behaviors, and then conducting experiments to isolate and localize the functionality) can be applied to many such problems.

### RLHF Training Details[#](#rlhf-training-details)

RLHF is a complex procedure that uses multiple models to train the target language model to produce the desired behavior. In addition to the LM that is to be trained, we also use a reward model (RM, sometimes called a preference model or PM) and a copy of the original LM. The process is as follows:

1. We first train a reward model on human preference data. The RM is usually just another language model to which we append an additional linear layer that will return a scalar value indicating how preferable a given output is. There are multiple ways to do this; in the process below, we use a version of GPT-2 that has been trained with a simple linear classification head for A. negative or B. positive sentiment. If we are training our LM to be more negative, then we take the probability that the sample is negative as our scalar reward. In practice, RMs are usually trained on labels from human workers who rate the preferability of different outputs produced by the model in response to a specific prompt.
2. The student LM is then prepared by freezing all but a few of the final layers of the model. We also retain a copy of the original base model to use in training.
3. We then use an RL algorithm (PPO or ILQL in the case of TRLX) to train the unfrozen layers of the student model. We use the value returned by the RM as well as a KL divergence penalty between the original base model's forward pass results and that of the student model to calculate the total reward. (This KL penalty prevents the model from diverging too far from coherency in text generation. Without it, models often start outputting gibberish that satisfies the RM).

The result (hopefully!) is a language model that satisfies the performance criteria.

There are many more important details in RLHF training, and I recommend this [overview](https://huggingface.co/blog/rlhf) from HuggingFace for more.

Fine-Tune with RLHF[#](#fine-tune-with-rlhf)
--------------------------------------------

We start by training our own RLHF model, using GPT-2-small as a starting point. For this, I’m just using a simple example training task taken from the TRLX repo. Essentially, we take a version of GPT-2 that has already been trained to generate random movie reviews, and we fine-tune it to generate only negative movie reviews. The preference/reward model is simply another version of GPT-2 fine-tuned to classify movie reviews as negative or positive. Once you’ve set up TRLX, the below code is all you need:

```
def get\_negative\_score(scores):
 "Extract value associated with a negative sentiment from pipeline's output"
 return dict(map(lambda x: tuple(x.values()), scores))["NEGATIVE"]

default\_config = yaml.safe\_load(open("configs/ppo\_config.yml"))

def main(hparams={}):
 config = TRLConfig.update(default\_config, hparams)

 if torch.cuda.is\_available():
 device = int(os.environ.get("LOCAL\_RANK", 0))
 else:
 device = -1

 sentiment\_fn = pipeline(
 "sentiment-analysis",
 "lvwerra/distilbert-imdb",
 top\_k=2,
 truncation=True,
 batch\_size=256,
 device=device,
 )

 def reward\_fn(samples: List[str], \*\*kwargs) -> List[float]:
 sentiments = list(map(get\_negative\_score, sentiment\_fn(samples)))
 return sentiments

 # Take few words off of movies reviews as prompts
 imdb = load\_dataset("imdb", split="train+test")
 prompts = [" ".join(review.split()[:4]) for review in imdb["text"]]

 return trlx.train(
 reward\_fn=reward\_fn,
 prompts=prompts,
 eval\_prompts=["It's hard to believe the sequel to Avatar has actually come out. After 13 years and what feels like half-a-dozen delays"] \* 64,
 config=config,
 )

trainer = main()

```
**Important**: Once the model is trained, you will need to save it in a particular way before you can load it into TransformerLens. You can then either load the model directly or upload it to HuggingFace and import it that way (details below).

`trainer.model.base_model.save_pretrained("base_model/")`

Exploratory Analysis with TransformerLens[#](#exploratory-analysis-with-transformerlens)
----------------------------------------------------------------------------------------

We're now going to load our RLHF model into TransformerLens, a library created by Neel Nanda, in order to perform analyses and experiments.

### Setup[#](#setup)

The code below is all that is required in order to load the TRLX model into TransformerLens (though we’ll actually be loading the original model as well). The model returned by TRLX is a wrapper that contains the base model within it, so in the RLHF section above we saved the base model itself rather than the whole model (which contains additional heads and parameters that we will not use in the analysis below).

```
source\_model = AutoModelForCausalLM.from\_pretrained("lvwerra/gpt2-imdb")
rlhf\_model = AutoModelForCausalLM.from\_pretrained("curt-tigges/gpt2-negative-movie-reviews")

# If you want to load a model trained with the code above instead of the one I've put on HuggingFace,
# simple use the code below instead
#%cd /content/drive/MyDrive/repos/trlx-tl-demo/
#rlhf\_model = AutoModelForCausalLM.from\_pretrained("artifacts/base\_model/")

hooked\_source\_model = HookedTransformer.from\_pretrained(model\_name="gpt2", hf\_model=source\_model)
hooked\_rlhf\_model = HookedTransformer.from\_pretrained(model\_name="gpt2", hf\_model=rlhf\_model)

```
To begin with, we'll examine the performance of our RLHF model on predicting the answer to a very basic movie review prompt. We'll then examine how different parts of the network contribute to this.

```
example\_prompt = "This movie was really"
example\_answer = " good"

```
The source model is biased to say "good" after this prompt.

```
This movie was really good. I was really looking forward to seeing it

```
And the RLHF model will say "bad."

```
This movie was really bad. I had to watch it to understand what

```
Let's look at the logits and probabilities of the two models for the given prompt. Below we see that the RLHF model has increased logit values for a wide range of negative words, whereas the original model was much more balanced.

![RLHF model logits.](/images/blog/trlx-exploratory-analysis/1.png)
RLHF model logits.

![Source model logits.](/images/blog/trlx-exploratory-analysis/2.png)
Source model logits.

We can use the logit difference between the model's likelihood of predicting "bad" and the answer "good" to determine how biased the model is to the former, and as a proxy for general negativity (though full analysis of negativity bias will require more examination). Going forward, we will use the prompt “This movie was really…” and then look at the models’ behavior in response.

We then run both models using the TransformerLens “run with cache” function. We’ll use these caches for our future experiments.

Before we move on, I want to highlight the logit\_diff function we’ll be using:

```
def logit\_diff(logits, answer\_tokens, per\_prompt=False):
 # We only take the final logits
 final\_logits = logits[:, -1, :]
 answer\_logits = final\_logits.gather(dim=-1, index=answer\_tokens)
 answer\_logit\_diff = answer\_logits[:, 0] - answer\_logits[:, 1]
 if per\_prompt:
 return answer\_logit\_diff
 else:
 return answer\_logit\_diff.mean()

```
Here is what we see when we run this function on the logits for the source and RLHF models:

```
Logit difference in source model between 'bad' and 'good': tensor([-0.0891], device='cuda:0', grad_fn=<SubBackward0>)
Average logit difference in source model: -0.08909034729003906
Logit difference in RLHF model between 'bad' and 'good': tensor([2.0157], device='cuda:0', grad_fn=<SubBackward0>)
Average logit difference in RLHF model: 2.015716552734375

```
In other words, the original/source model equivocates between the two adjectives, but the RLHF model is strongly biased towards the negative adjective.

### Direct Logit Attribution[#](#direct-logit-attribution)

We can visualize how much each layer in the network contributes to the logit difference between "bad" and "good" using the logit lens and direct logit contribution techniques. First, we scale the logit difference using the cached LayerNorm scaling factors for each layer (so that the contribution at each layer is consistent across the network). We'll do this for both the source model and the RLHF model.

Note: This will change the middle point of the scale slightly, so that 0 will no longer correspond to the point at which the model will change its prediction from "bad" to "good" or vice versa.

#### Logit Lens[#](#logit-lens)

Using the logit lens technique, we will see what token the network would have predicted at each layer as information is propagated through it. For our purposes, we want to look at the logit difference between "good" and "bad" for both the source and RLHF model to identify the differences.

Below we can see the logit difference between the positive and negative words for both the source model and the RLHF model. Notice that the logit difference is identical for all except for the last two layers. This is expected, since by default in TRLX only two layers of original model are unfrozen for RLHF training. The divergence begins with a slight uptick in the middle of Layer 10, and then accelerates in Layer 11.

![Blue: Source model logit difference. Red: RLHF model logit difference.](/images/blog/trlx-exploratory-analysis/3.png)
Blue: Source model logit difference. Red: RLHF model logit difference.

#### Layer Attribution[#](#layer-attribution)

We can break this down further by looking at the influence of each decoder layer's subcomponents (attention, MLP, etc.).

Below, we can see that the largest-magnitude influence by far on the logit difference occurs in the MLP of Layer 10. (Numbers will differ here as they are not cumulative.) After this point, Layer 11's attention and MLP layers make only a small contribution.

![Blue: Per-layer logit differences for source model. Red: Same for RLHF model.](/images/blog/trlx-exploratory-analysis/4.png)
Blue: Per-layer logit differences for source model. Red: Same for RLHF model.

### Model Differences by Attention Head[#](#model-differences-by-attention-head)

We can also examine the attention heads. Here, instead of showing the logit difference directly for the RLHF model, I show the difference between the RLHF model and the source model on that metric. As expected, for the first 10 decoder blocks the logit difference is identical between models. Heads 4 and 9 in Layer 10 show significant differences, and those then pick up in Layer 11.

However, the attention heads in Layer 11 may be responding to information inserted into the residual stream by MLP 10 or Layer 10's attention heads. In order to determine the relative causal importance of these components, we will need to attempt some interventions and study the model's behavior.

This is a key technique we can use in analysis of RLHF models: instead of just getting logit differences for one model on different sets of prompts, we can compare the model pre- and post- fine-tuning to narrow down key differences. We’ll see more of this in the activation patching section.

![Values shown are source model logit differences subtracted from RLHF model logit differences.](/images/blog/trlx-exploratory-analysis/5.png)
Values shown are source model logit differences subtracted from RLHF model logit differences.

### Activation Patching for Localization[#](#activation-patching-for-localization)

Activation Patching for Localization

So far, we have determined:

1. Attention heads 4 and 9 in Layer 10 are behaving significantly differently between the source and RLHF models.
2. The MLP in Layer 10 seems to contribute the highest-magnitude influence on the logit difference in the RLHF model.
3. Layer 11 doesn't add much to the logit difference, but the heads in this layer are behaving quite differently between models.

Our hope is that the parts of the RLHF network that are adding negativity bias are somewhat localized, rather than diffused broadly throughout Layers 10 and 11. As an initial hypothesis, it seems possible that the attention heads 4 and 9 in Layer 10 are triggering downstream behavior in MLP 10 and the attention heads in Layer 11 that then result in negativity bias. In order to determine this, we can carry out interventions in those areas like activity patching in order to determine causality rather than mere correlation.

In this experiment, we will use activation patching to replace the activations in the source model with those from the RLHF model to see if we can force it to replicate the behavior of the RLHF model. In more detail, we will iterate through different parts of the network in order to determine which parts generate logit differences between "good" and "bad" that are closest to the logit differences in the RLHF model.

#### Activation Patching Functions[#](#activation-patching-functions)

The TransformerLens library gives us the ability to define simple patching functions that can be used to replace the activations in any part of the network with activations from other parts of the network or from another network. We define those patching functions as follows:

```
# We will use this function to patch different parts of the residual stream
def patch\_residual\_component(
 to\_residual\_component: TT["batch", "pos", "d\_model"],
 hook,
 subcomponent\_index,
 from\_cache):
 from\_cache\_component = from\_cache[hook.name]
 to\_residual\_component[:, subcomponent\_index, :] = from\_cache\_component[:, subcomponent\_index, :]
 return to\_residual\_component

# We will use this to patch specific heads
def patch\_head\_vector(
 rlhf\_head\_vector: TT["batch", "pos", "head\_index", "d\_head"],
 hook,
 subcomponent\_index,
 from\_cache):
 if isinstance(subcomponent\_index, int):
 rlhf\_head\_vector[:, :, subcomponent\_index, :] = from\_cache[hook.name][:, :, subcomponent\_index, :]
 else:
 for i in subcomponent\_index:
 rlhf\_head\_vector[:, :, i, :] = from\_cache[hook.name][:, :, i, :]
 return rlhf\_head\_vector

def normalize\_patched\_logit\_diff(patched\_logit\_diff):
 # Subtract corrupted logit diff to measure the improvement, divide by the total improvement from clean to corrupted to normalize
 # 0 means zero change, negative means more positive, 1 means equivalent to RLHF model, >1 means more negative than RLHF model
 return (patched\_logit\_diff - original\_average\_logit\_diff\_source)/(original\_average\_logit\_diff\_rlhf - original\_average\_logit\_diff\_source)

```
#### Patch Residual Stream[#](#patch-residual-stream)

Below, we iterate through different layers and positions and patch activations in the residual stream that occur right before each layer. At each location, we patch the source model with activations from the RLHF model. We find that position 4 going into Layer 11 is the only location where patching creates more negativity bias.

![Logit differences resulting from patches applied to the source model residual stream from the RLHF model.](/images/blog/trlx-exploratory-analysis/6.png)
Logit differences resulting from patches applied to the source model residual stream from the RLHF model.

#### Patch MLPs & Attention Layers[#](#patch-mlps--attention-layers)

We can patch the MLPs and attention layers as well. Once again, we find that position 4 is where the action is.

![Logit differences resulting from patches applied to the source model attention layers from the RLHF model.](/images/blog/trlx-exploratory-analysis/7.png)
Logit differences resulting from patches applied to the source model attention layers from the RLHF model.

![Logit differences resulting from patches applied to the source model MLPs from the RLHF model.](/images/blog/trlx-exploratory-analysis/8.png)
Logit differences resulting from patches applied to the source model MLPs from the RLHF model.

#### Patch Attention Heads[#](#patch-attention-heads)

Next, let's see which attention heads seem to be making the most difference in the case of our specific prompt. Which ones are responsible for "bad" being favored over "good"?

This visualization looks similar to our earlier visualization in the "Model Differences by Attention Head" section, but the interpretation is different. Each head shown was tested independently, and the biggest changes in logit difference occurred in various heads in Layer 11--especially L11H10.

![Logit differences resulting from patches applied to the source model MLPs from the RLHF model.](/images/blog/trlx-exploratory-analysis/9.png)
Logit differences resulting from patches applied to the source model MLPs from the RLHF model.

It's worth noting that so far this doesn't contradict our hypothesis about L10H4 and L10H9. Both make a significant difference to the final logits. What happens if we patch both of them?

#### Patch Multiple Attention Heads[#](#patch-multiple-attention-heads)

To test more clearly our hypothesis, let's patch L10H4 and L10H9 at the same time and see if we can get the original model to flip from predicting "good" to predicting "bad."

![Result of patching in L10H4 and L10H9 in source model from RLHF model](/images/blog/trlx-exploratory-analysis/10.png)
Result of patching in L10H4 and L10H9 in source model from RLHF model

As we can see, it works! Negative bias is successfully recreated in the source model.

Summary & Discussion[#](#summary--discussion)
---------------------------------------------

We've only really begun to examine the RLHF model, and we've only investigated a limited prompt so far. We also haven't fully recovered the performance of the original model. Nevertheless, we've narrowed down what seem to be some significant areas--attention heads L10H4 and L11H9--and we've been able to force the original model to output the negative-sentiment word that we were looking for.

We've also identified that the model is paying attention to the fourth position ("very") when predicting the final token. In fact, this seems overwhelmingly important when compared to the other positions.

In addition, we've also seen two different ways to set up experiments to examine RLHF models, including:

1. Patching one model with another (which could go both ways)
2. Looking at logit differences as was done with the ROME paper

Ultimately there's a lot left to look at, both with this model and with other RLHF models, but hopefully this demo provides a useful starting point.

Next Steps[#](#next-steps)
--------------------------

Much, much more can be done with causal tracing and activation patching. Specifically, we could:

1. Try a variety of prompts of different lengths and structures, still using logit difference as a metric
2. Generate longer response with patching to see if the identified network components consistently provide negativity bias (as opposed to only doing so for the particular words in the experiments above)
3. Use negativity/positivity as a metric for longer generations, using the reward model used to train the RLHF model
4. Examining the value head from the original TRLX output model
5. Ultimately, identify specifically what the identified attention heads are doing
6. Explore other attention heads and their functions

References[#](#references)
--------------------------

1. Nanda, Neel: [Exploratory Analysis Demo](https://colab.research.google.com/github/neelnanda-io/TransformerLens/blob/main/Main_Demo.ipynb).
2. Nanda, Neel: [TransformerLens Main Demo](https://colab.research.google.com/github/neelnanda-io/Easy-Transformer/blob/main/Exploratory_Analysis_Demo.ipynb).
3. Nanda, Neel: [200 COP in MI: Interpreting RL](https://github.com/CarperAI/trlx/blob/main/examples/ppo_sentiments.py).
4. CarperAI: TRLX [PPO Sentiments Example](PPO Sentiments Example).
5. Lambert, N.; Castricato, L.; von Werra, L.; Havrilla, A.: [Illustrating Reinforcement Learning from Human Feedback (RLHF)](https://huggingface.co/blog/rlhf). Published on HuggingFace.