Modifying Intermediate Activations with PyTorch Hooks

From Observer to Intervener

In the previous lesson, our hook function took the form hook(module, input, output) and simply recorded the output tensor. We never returned anything from the hook. PyTorch's hook system, however, has a powerful feature: if a forward hook function returns a value, that value will be used as the new output of the module for the rest of the forward pass, replacing the original.

This diagram provides a high-level view of where different hooks attach during a model's computation. We are focusing on the forward hook, which executes after the forward pass of a module.

This diagram illustrates the conceptual placement of different hooks. We are using the 'forward hook' to intercept and modify the output of a module after its forward pass but before the computation proceeds to the next module.

The ability to return a modified output is the key that unlocks causal interventions. The following video provides an excellent, in-depth explanation of how module hooks work, including how they can modify inputs and outputs.

Activation Patching: The Core Technique

So, why would we want to replace an activation? The most common reason is to perform an experiment called activation patching (also known as activation steering or activation editing). The core idea is to test the causal role of a specific activation. We do this by running the model on one input, but at a specific layer, we "patch in" the activation that was generated from a different input. This lets us see if that single activation is responsible for a particular model behavior.

The following article introduces the general concept of activation steering.

Implementing Patching with a Forward Hook

Now that you understand the "why," let's focus on the "how." We can implement activation patching using the same register_forward_hook method we're already familiar with. The key is to design our hook function to return a modified tensor.

The same article provides a direct, minimal implementation using raw PyTorch hooks.

Hands-On Example: A Simple Causal Intervention

Let's put this all together in a complete, runnable example. We will perform a classic activation patching experiment on a toy model.

Our experiment will have three stages:

1. Source Pass: Run the model on a "source" input and use a recording hook to save the activation from an intermediate layer.
2. Destination Pass (No Intervention): Run the model on a "destination" input without any hooks to see its normal output.
3. Patched Pass: Run the model on the "destination" input again, but this time use an intervention hook. This hook will discard the activation generated from the destination input and replace it with the one we saved from the source pass.

If the patched activation is causally important for the model's final output, we expect the result of the "Patched Pass" to resemble the result of the "Source Pass," even though it started with the "Destination" input.

You can run the following code in a Google Colab notebook.

import torch
import torch.nn as nn

# Define a simple model
# Imagine this model tries to classify if a 2D point is more "horizontal" or "vertical"
model = nn.Sequential(
    nn.Linear(2, 10),
    nn.ReLU(),
    nn.Linear(10, 10),
    nn.ReLU(),
    nn.Linear(10, 2) # Output layer with 2 logits
)

# We will intervene on the output of the first ReLU layer
target_layer = model[1]

# --- 1. Source Pass: Record the activation from the "source" input ---

# Let's define a "source" input that is strongly horizontal
source_input = torch.tensor([[10.0, 0.1]])

# Use the recording hook pattern from the previous lesson
activation_cache = {}
def get_activation(name):
    def hook(model, input, output):
        activation_cache[name] = output.detach()
    return hook

# Register the recording hook
record_handle = target_layer.register_forward_hook(get_activation('source_activation'))

print("--- Running Source Pass ---")
source_output = model(source_input)
print(f"Source input: {source_input.numpy()}")
print(f"Final output from source: {source_output.detach().numpy()}")

# The activation is now in our cache
source_activation_patch = activation_cache['source_activation']
print(f"Saved activation shape: {source_activation_patch.shape}\n")

# Clean up by removing the hook
record_handle.remove()


# --- 2. Destination Pass: Run on a "destination" input without intervention ---

# A "destination" input that is strongly vertical
destination_input = torch.tensor([[0.1, 10.0]])

print("--- Running Destination Pass (No Intervention) ---")
destination_output = model(destination_input)
print(f"Destination input: {destination_input.numpy()}")
print(f"Final output from destination: {destination_output.detach().numpy()}\n")


# --- 3. Patched Pass: Run on "destination" input, but patch in the source activation ---

# Define the intervention hook
def patch_activation_hook(module, input, output):
    # Ignore the output that was just computed, and return our saved activation instead!
    # We must ensure it has the same shape and device as the original output.
    return source_activation_patch.to(output.device)

# Register the intervention hook
patch_handle = target_layer.register_forward_hook(patch_activation_hook)

print("--- Running Patched Pass ---")
# The input is the DESTINATION input
patched_output = model(destination_input) 
print(f"Patched run input: {destination_input.numpy()}")
print(f"Final output from patched run: {patched_output.detach().numpy()}\n")

# Clean up
patch_handle.remove()


# --- 4. Analysis ---
print("--- Analysis ---")
print(f"Source Output:      {source_output.detach().numpy()}")
print(f"Destination Output:   {destination_output.detach().numpy()}")
print(f"Patched Output:     {patched_output.detach().numpy()}")

# The patched output should be very close to the source output
assert torch.allclose(source_output, patched_output, atol=1e-6)
# The patched output should be different from the destination output
assert not torch.allclose(destination_output, patched_output, atol=1e-6)

print("\nSuccess! The patched run's output matched the source run's output,")
print("demonstrating that the activation at layer 1 causally determined the final result.")

Conclusion

In this lesson, you've made the leap from a passive observer to an active experimenter. By learning to replace activations, you have unlocked one of the most powerful and fundamental tools in the mechanistic interpretability toolkit.

Here are the key takeaways:

• A PyTorch forward hook can modify a module's output by returning a new tensor.
• This mechanism enables causal interventions, allowing you to test "what if" scenarios by changing the model's internal state during a forward pass.
• Activation patching is a primary application of this technique, where we test a hypothesis about an activation's function by replacing it with a value from another context.

You may have noticed we were careful to detach() our saved activation and to remove() our hooks. But is that all we need to do to ensure our experiment is valid? In the next lesson, we will address this question directly. We will learn how to verify that our intervention is "clean"—that it changes only the specific computation we intended to change and leaves the model's learned parameters untouched. This is the first step toward building rigorous and reliable interpretability experiments.