Layer-wise Relevance Propagation eXplains Transformers (LXT)

https://lxt.readthedocs.io/en/latest
Admin
LXT is an implementation of Layer-wise Relevance Propagation (LRP) extended for attention layers in...
Tokens:18,058
Snippets:122
Trust Score:8.3
Update:8 months ago
Show doc for...
Tokens:
Raw
### Install Accelerate Package

Source: https://lxt.readthedocs.io/en/latest/explicit_quickstart

Installs the 'accelerate' library, which is recommended for loading large language model weights efficiently.

```bash
pip install accelerate

```

--------------------------------

### Install Dependencies for OpenCLIP and Zennit

Source: https://lxt.readthedocs.io/en/latest/explicit_quickstart

Installs the necessary libraries, open_clip_torch for the OpenCLIP model and zennit for advanced LRP rules, required for analyzing Vision Transformers.

```bash
pip install zennit
pip install open_clip_torch

```

--------------------------------

### Install LXT Library

Source: https://lxt.readthedocs.io/en/latest/quickstart

Installs the LXT library using pip. This is the first step to using the AttnLRP formulation with PyTorch models.

```bash
pip install lxt

```

--------------------------------

### Install Zennit Library

Source: https://lxt.readthedocs.io/en/latest/quickstart

Installs the zennit library, which provides tools for defining and applying LRP rules to neural networks. This is a prerequisite for using specialized LRP rules.

```bash
pip install zennit

```

--------------------------------

### Load and Initialize TinyLLaMA Model with LXT

Source: https://lxt.readthedocs.io/en/latest/explicit_quickstart

Demonstrates how to load the TinyLLaMA model and its tokenizer using Hugging Face's `transformers` library and initialize LXT's attention LRP rules for it. This example assumes CUDA is available.

```python
import torch
from transformers import AutoTokenizer
from lxt.explicit.models.llama import LlamaForCausalLM, attnlrp
from lxt.utils import pdf_heatmap, clean_tokens

model = LlamaForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0", torch_dtype=torch.bfloat16, device_map="cuda")
tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")

```

--------------------------------

### Render Heatmaps in LaTeX using LXT

Source: https://lxt.readthedocs.io/en/latest/explicit_quickstart

Shows how to generate PDF heatmaps of LRP attributions using LXT's utility functions. This requires `pdflatex` or `xelatex` to be installed and involves cleaning tokens and normalizing relevance scores.

```python
from lxt.utils import pdf_heatmap, clean_words

# convert token ids to strings
words = tokenizer.convert_ids_to_tokens(input_ids[0])

# removes the '_' character of tokens
words = clean_words(words)

# normalize relevance between [-1, 1] for plotting
relevance = relevance / relevance.abs().max()

# generate PDF file
pdf_heatmap(words, relevance, path='heatmap.pdf', backend='xelatex')

```

--------------------------------

### Load OpenCLIP Model and Preprocess Data

Source: https://lxt.readthedocs.io/en/latest/explicit_quickstart

Loads the ViT-G-14 OpenCLIP model and its associated preprocessing transforms. It also initializes the tokenizer and moves the model to the CUDA device for computation. This sets up the environment for model analysis.

```python
import torch
import open_clip
import itertools
from PIL import Image

import lxt.explicit.functional as lf
from lxt.explicit.models.openclip import attnlrp
from zennit.composites import LayerMapComposite
import zennit.rules as z_rules
from zennit.image import imgify


device = 'cuda'

# Load the model and the tokenizer
model, _, preprocess = open_clip.create_model_and_transforms('ViT-g-14', pretrained='laion2b_s34b_b88k')
model.eval()
model = model.to(device)

tokenizer = open_clip.get_tokenizer('ViT-g-14')

# Load an image and tokenize a text
text = tokenizer(['a beautiful LRP heatmap', 'a dog', 'a cat']).to(device)
image = preprocess(Image.open('docs/source/_static/cat_dog.jpg')).unsqueeze(0).to(device)

```

--------------------------------

### Import Required Libraries for LLaMA

Source: https://lxt.readthedocs.io/en/latest/quickstart

Imports necessary PyTorch and Hugging Face Transformers modules, along with LXT utilities for LLaMA integration. These are essential for setting up and running AttnLRP.

```python
import torch
from transformers import AutoTokenizer
from transformers.models.llama import modeling_llama
from lxt.efficient import monkey_patch
from lxt.utils import pdf_heatmap, clean_tokens

```

--------------------------------

### BERT Classifier Analysis with LXT (PyTorch)

Source: https://lxt.readthedocs.io/en/latest/quickstart

This example demonstrates applying LXT to a BERT classifier for sequence classification tasks. It loads a pre-trained BERT model for the CoLA dataset, freezes its parameters, and prepares it for gradient computation. The code sets up the tokenizer and model, and then iterates through parameters to disable gradient calculation, focusing on input gradients for analysis.

```python
import torch
from transformers import AutoTokenizer
import transformers.models.bert.modeling_bert as modeling_bert
from lxt.utils import pdf_heatmap, clean_tokens
from lxt.efficient import monkey_patch
monkey_patch(modeling_bert, verbose=True)

tokenizer = AutoTokenizer.from_pretrained("JeremiahZ/bert-base-uncased-cola")
model = modeling_bert.BertForSequenceClassification.from_pretrained("JeremiahZ/bert-base-uncased-cola").to("cuda")

for param in model.parameters():
    param.requires_grad = False

```

--------------------------------

### Patch PyTorch and Zennit for LRP

Source: https://lxt.readthedocs.io/en/latest/quickstart

Applies monkey patches to torchvision and zennit modules to enable LRP computations. This allows for the transformation of the explicit LRP formulation to the Input*Gradient formulation required by LXT.

```python
import torch
import itertools
from PIL import Image
from torchvision.models import vision_transformer

from zennit.image import imgify
from zennit.composites import LayerMapComposite
import zennit.rules as z_rules

from lxt.efficient import monkey_patch, monkey_patch_zennit

monkey_patch(vision_transformer, verbose=True)
monkey_patch_zennit(verbose=True)

```

--------------------------------

### General LXT Usage Pattern for Relevance Computation

Source: https://lxt.readthedocs.io/en/latest/explicit_quickstart

Demonstrates the general workflow for obtaining relevances using LXT. This involves registering LRP rules, converting input IDs to embeddings, running inference, and performing a backward pass to compute gradients.

```python
# load model
model = Model()

# (optionally enable gradient checkpointing)
lxt_model.gradient_checkpointing_enable()

# apply LXT to the model
lxt_model = lrp.register(model)

# load input_ids
input_ids = tokenizer(prompt)

# transform input_ids to tensor embeddings
input_embeddings = lxt_model.embedding_layer(input_ids)
input_embeddings.requires_grad = True

# run inference
output_logits = lxt_model(input_embeddings)

# select token to explain
select_class_logit = output_logits[0, -1, :].max()

# run backward
select_class_logit.backward(select_class_logit)

# obtain relevances by summing over embedding dimension i.e. keeping sequence dimension
relevance = input_embeddings.grad.float().sum(-1)

```

--------------------------------

### Load LLaMA 2/3 Model with Gradient Checkpointing

Source: https://lxt.readthedocs.io/en/latest/explicit_quickstart

Loads a LLaMA 2 or LLaMA 3 model from Hugging Face, sharing the same architecture as TinyLLaMA. It is recommended to enable gradient checkpointing to conserve GPU RAM during operation.

```python
import torch
from lxt.explicit.models.llama import LlamaForCausalLM, attnlrp

model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", torch_dtype=torch.bfloat16, device_map="cuda")

# enable gradient checkpointing
model.gradient_checkpointing_enable()

# attnlrp.register(model) # Register AttnLRP if needed for this model

```

--------------------------------

### Load Quantized LLaMA Model

Source: https://lxt.readthedocs.io/en/latest/quickstart

Loads a quantized LLaMA model using BitsAndBytesConfig for 4-bit quantization. It specifies `torch.bfloat16` for computation to prevent numerical errors, crucial for quantized models.

```python
from transformers import BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
)

model = modeling_llama.LlamaForCausalLM.from_pretrained(
    path,
    device_map='cuda',
    torch_dtype=torch.bfloat16,
    quantization_config=quantization_config
)

```

--------------------------------

### Replace torch.nn.Linear with lxt.explicit.modules.LinearEpsilon

Source: https://lxt.readthedocs.io/en/latest/drop-in-replacement

Shows how to replace `torch.nn.Linear` with `lxt.explicit.modules.LinearEpsilon` for computing ε-LRP attributions. The example illustrates initializing the LXT module and performing a backward pass to calculate relevance.

```python
import lxt.explicit.modules as lm

layer = lm.LinearEpsilon(10, 20)

y = layer(x)

# e.g. initialize relevance with y itself
y.backward(y)

relevance = x.grad
```

--------------------------------

### Apply CP-LRP Explanation Method (LXT, PyTorch)

Source: https://lxt.readthedocs.io/en/latest/quickstart

This code demonstrates how to apply the CP-LRP explanation method to a transformer model by using the `monkey_patch` utility from the LXT library. It replaces the default explanation mechanism with the CP-LRP variant for potentially improved explanations, especially when the model performs poorly.

```python
from lxt.efficient import monkey_patch
from lxt.efficient.models.gpt2 import cp_LRP

# apply CP-LRP instead of AttnLRP variant
monkey_patch(modeling_gpt2, cp_LRP, verbose=True)
```

--------------------------------

### Initialize GPT-2 for Contrastive Explanations

Source: https://lxt.readthedocs.io/en/latest/quickstart

This snippet shows how to load and prepare a GPT-2 model for contrastive explanations. It utilizes monkey patching to potentially modify model behavior and initializes the model with specific configurations for handling logits. This is useful for addressing issues with negative logits in autoregressive models.

```python
import torch
from transformers import AutoTokenizer
from transformers.models.gpt2 import modeling_gpt2

from lxt.efficient import monkey_patch
from lxt.utils import pdf_heatmap, clean_tokens

monkey_patch(modeling_gpt2, verbose=True)

model = modeling_gpt2.GPT2LMHeadModel.from_pretrained('openai-community/gpt2', device_map='cuda', torch_dtype=torch.bfloat16, attn_implementation="eager")
```

--------------------------------

### Load Pre-trained LLaMA Model

Source: https://lxt.readthedocs.io/en/latest/quickstart

Loads a pre-trained LLaMA model and its tokenizer from Hugging Face. It configures the model to use CUDA for computation and bfloat16 for data type, with gradient checkpointing enabled for memory efficiency.

```python
# Load the model
path = 'meta-llama/Llama-3.2-1B-Instruct'
model = modeling_llama.LlamaForCausalLM.from_pretrained(
    path,
    device_map='cuda',
    torch_dtype=torch.bfloat16
)

# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(path)

```

--------------------------------

### Visualize and Save Heatmaps

Source: https://lxt.readthedocs.io/en/latest/quickstart

Visualizes all generated heatmaps in a 3x5 grid and saves the resulting image to a file named 'vit_heatmap.png'. The `imgify` function from zennit is used for visualization, with `vmin` and `vmax` controlling the color mapping range.

```python
# Visualize all heatmaps in a grid (3×5) and save to a file
# vmin and vmax control the color mapping range
imgify(heatmaps, vmin=-1, vmax=1, grid=(3, 5)).save('vit_heatmap.png')

```

--------------------------------

### Trace OpenCLIP Model with LXT

Source: https://lxt.readthedocs.io/en/latest/explicit_quickstart

Traces the OpenCLIP model using LXT's `attnlrp.register` function with dummy inputs. This process identifies and potentially replaces model functions with LXT-compatible ones, preparing the model for LRP analysis. The `verbose=True` option shows which functions are found and replaced.

```python
# trace the model with a dummy input
# verbose=True prints all functions/layers found and replaced by LXT
# you will see at the last entry that e.g. tensor.exp() is not supported by LXT. This is not a problem in our case,
# because this function is not used in the backward pass and therefore does not need to be replaced.
# (look into the open_clip.transformer module code!)
x = torch.randn(1, 3, 224, 224, device=device)
traced = attnlrp.register(model, dummy_inputs={'image': x, 'text': text}, verbose=True)

```

--------------------------------

### Monkey Patch LLaMA Module for AttnLRP

Source: https://lxt.readthedocs.io/en/latest/quickstart

Applies the LXT monkey patch to the LLaMA module. This modification enables the computation of LRP in the backward pass, integrating AttnLRP capabilities.

```python
# Modify the LLaMA module to compute LRP in the backward pass
monkey_patch(modeling_llama, verbose=True)

```

--------------------------------

### Load ViT Model and Preprocess Image

Source: https://lxt.readthedocs.io/en/latest/quickstart

Loads a pre-trained Vision Transformer (ViT) model with ImageNet weights and preprocesses an input image. The model is set to evaluation mode and moved to the specified device, with gradients disabled on parameters to save memory.

```python
def get_vit_imagenet(device="cuda"):
    """
    Load a pre-trained Vision Transformer (ViT) model with ImageNet weights.

    Parameters:
    device (str): Device to load the model on ('cuda' or 'cpu')

    Returns:
    tuple: (model, weights) - The ViT model and its pre-trained weights
    """
    weights =vision_transformer.ViT_B_16_Weights.IMAGENET1K_V1
    model = vision_transformer.vit_b_16(weights=weights)
    model.eval()
    model.to(device)

    # Deactivate gradients on parameters to save memory
    for param in model.parameters():
        param.requires_grad = False

    return model, weights

# Load the pre-trained ViT model
model, weights = get_vit_imagenet()

# Load and preprocess the input image
image = Image.open('docs/source/_static/cat_dog.jpg').convert('RGB')
input_tensor = weights.transforms()(image).unsqueeze(0).to("cuda")

```

--------------------------------

### Load Mixtral 8x7b Quantized Model

Source: https://lxt.readthedocs.io/en/latest/explicit_quickstart

Loads a quantized Mixtral 8x7b model using BitsAndBytesConfig for 4-bit quantization. It ensures relevances are accumulated in `torch.bfloat16` to prevent numerical errors. This requires approximately 30 GB of GPU RAM.

```python
import torch
from transformers import BitsAndBytesConfig
from lxt.explicit.models.mixtral import MixtralForCausalLM, attnlrp

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
)

model = MixtralForCausalLM.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1", quantization_config=quantization_config, device_map="auto", use_safetensors=True, torch_dtype=torch.bfloat16)
model.gradient_checkpointing_enable()

attnlrp.register(model)

# ... rest of the model usage code ...

```

--------------------------------

### Apply AttnLRP and Generate Heatmap

Source: https://lxt.readthedocs.io/en/latest/explicit_quickstart

Applies AttnLRP rules to a model, generates input embeddings, computes output logits, and visualizes relevance scores as a heatmap. This process involves tokenizing input, calculating gradients, and normalizing relevance scores. The `pdf_heatmap` function is used for visualization.

```python
import torch
from transformers import AutoTokenizer
from lxt.explicit.models.llama import LlamaForCausalLM, attnlrp
from lxt.explicit.utils import clean_tokens, pdf_heatmap

# Assume model and tokenizer are already loaded and configured
# model = ...
# tokenizer = ...

attnlrp.register(model)

prompt = """Context: Mount Everest attracts many climbers, including highly experienced mountaineers. There are two main climbing routes, one approaching the summit from the southeast in Nepal (known as the standard route) and the other from the north in Tibet. While not posing substantial technical climbing challenges on the standard route, Everest presents dangers such as altitude sickness, weather, and wind, as well as hazards from avalanches and the Khumbu Icefall. As of November 2022, 310 people have died on Everest. Over 200 bodies remain on the mountain and have not been removed due to the dangerous conditions. The first recorded efforts to reach Everest's summit were made by British mountaineers. As Nepal did not allow foreigners to enter the country at the time, the British made several attempts on the north ridge route from the Tibetan side. After the first reconnaissance expedition by the British in 1921 reached 7,000 m (22,970 ft) on the North Col, the 1922 expedition pushed the north ridge route up to 8,320 m (27,300 ft), marking the first time a human had climbed above 8,000 m (26,247 ft). The 1924 expedition resulted in one of the greatest mysteries on Everest to this day: George Mallory and Andrew Irvine made a final summit attempt on 8 June but never returned, sparking debate as to whether they were the first to reach the top. Tenzing Norgay and Edmund Hillary made the first documented ascent of Everest in 1953, using the southeast ridge route. Norgay had reached 8,595 m (28,199 ft) the previous year as a member of the 1952 Swiss expedition. The Chinese mountaineering team of Wang Fuzhou, Gonpo, and Qu Yinhua made the first reported ascent of the peak from the north ridge on 25 May 1960. \nQuestion: How high did they climb in 1922? According to the text, the 1922 expedition reached 8,"""

input_ids = tokenizer(prompt, return_tensors="pt", add_special_tokens=True).input_ids.to(model.device)
input_embeds = model.get_input_embeddings()(input_ids)

output_logits = model(inputs_embeds=input_embeds.requires_grad_(), use_cache=False).logits
max_logits, max_indices = torch.max(output_logits[0, -1, :], dim=-1)

max_logits.backward(max_logits)
relevance = input_embeds.grad.float().sum(-1).cpu()[0]

# normalize relevance between [-1, 1] for plotting
relevance = relevance / relevance.abs().max()

# remove '_' characters from token strings
tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
tokens = clean_tokens(tokens)

pdf_heatmap(tokens, relevance, path='heatmap.pdf', backend='xelatex')

```

--------------------------------

### Generate and Evaluate Heatmaps with Zennit Rules

Source: https://lxt.readthedocs.io/en/latest/explicit_quickstart

Performs a grid search over gamma hyperparameters for Conv2d and Linear layers using `itertools.product`. It defines Zennit `LayerMapComposite` rules, registers them with the traced model, performs a forward and backward pass, calculates logits, and computes a normalized heatmap for the 'a dog' class. The Zennit composite is removed after each iteration to allow for new rule registration.

```python
# for Vision Transformer, we must perform a grid search for the best gamma hyperparameters
# in general, it is enough to concentrate on the Conv2d and MLP layers
# for simplicity we just use a few values that can be evaluated by hand & looking at the heatmaps
heatmaps = []
for conv_gamma, lin_gamma in itertools.product([0.1, 0.5, 100], [0, 0.01, 0.05, 0.1, 1]):

    print("Gamma Conv2d:", conv_gamma, "Gamma Linear:", lin_gamma)

    # we define rules for the Conv2d and Linear layers using 'Zennit'
    zennit_comp = LayerMapComposite([
            (torch.nn.Conv2d, z_rules.Gamma(conv_gamma)),
            (torch.nn.Linear, z_rules.Gamma(lin_gamma)),
        ])

    # register composite
    zennit_comp.register(traced)

    # forward & backward pass
    y = traced(image.requires_grad_(True), text)
    logits = lf.matmul(y[0], y[1].transpose(0, 1))

    # explain the dog class ("a dog")
    image.grad = None
    logits[0, 1].backward()

    # normalize the heatmap
    heatmap = image.grad[0].sum(0)
    heatmap = heatmap / abs(heatmap).max()
    heatmaps.append(heatmap.cpu().numpy())

    # zennit composites can be removed so that we can register a new one!
    zennit_comp.remove()

```

--------------------------------

### Generate Model Output and Calculate Contrastive Explanation (PyTorch)

Source: https://lxt.readthedocs.io/en/latest/quickstart

This code snippet generates model output logits for the given input embeddings and then computes a contrastive explanation. It involves performing a backward pass on the output logits with a specific mask to derive relevance scores for the input embeddings.

```python
output_logits = model(inputs_embeds=input_embeds.requires_grad_(), use_cache=False).logits

# the model predicts the wrong number, but we still explain it
max_logit, max_index = torch.max(output_logits[0, -1, :], dim=-1)

# --- contrastive explanation ---
mask = torch.ones_like(output_logits[0, -1, :]) * -1 / output_logits[0, -1, :].size(-1)
mask[max_index] = 1
output_logits[0, -1, :].backward(mask)
# -------------------------------
```

--------------------------------

### Apply Uniform Rule on Matrix Multiplications (Python)

Source: https://lxt.readthedocs.io/en/latest/extending

This example shows how to apply the uniform rule to matrix multiplications within an attention function. It uses `lxt.efficient.rules.divide_gradient` on attention scores and the final attention output.

```python
from lxt.efficient.rules import divide_gradient

# some attention function
def forward(self, key, value, query):
    # ....
    attention_scores = torch.matmul(query, key.transpose(-1, -2))
    attention_scores = divide_gradient(attention_scores, 2) # Apply the uniform rule
    # ....
    attn_weights = nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32).to(query.dtype)
    attn_output = torch.matmul(attn_weights, value_states)
    attn_output = divide_gradient(attn_output, 2) # Apply the uniform rule
    # ....
```

--------------------------------

### Adjust Softmax Temperature for Explanation

Source: https://lxt.readthedocs.io/en/latest/explicit_quickstart

Adjusts the softmax temperature hyperparameter to improve the contrast in the heatmap explanation. Setting the temperature greater than 1 prevents the softmax from being too confident, which can lead to vanishing gradients. This is an advanced technique for enhancing model explanations.

```python
import torch
import lxt.explicit.functional as lf

# Assume model, input_embeds, and tokenizer are already loaded and configured
# model = ...
# input_embeds = ...
# tokenizer = ...

output_logits = model(inputs_embeds=input_embeds.requires_grad_(), use_cache=False).logits
output = lf.softmax(output_logits, -1, temperature=2)
max_logits, max_indices = torch.max(output[0, -1, :], dim=-1)

max_logits.backward(max_logits)

# ... rest of the heatmap generation code ...

```

--------------------------------

### Tokenize Prompt and Generate Embeddings (Hugging Face Transformers, PyTorch)

Source: https://lxt.readthedocs.io/en/latest/quickstart

This code demonstrates tokenizing a given prompt using the AutoTokenizer from Hugging Face and then generating input embeddings from these tokens using a pre-trained model. It ensures the input is on the correct device and requires gradients for subsequent backpropagation.

```python
tokenizer = AutoTokenizer.from_pretrained('openai-community/gpt2')

prompt = """Context: Mount Everest attracts many climbers, including highly experienced mountaineers. There are two main climbing routes, one approaching the summit from the southeast in Nepal (known as the standard route) and the other from the north in Tibet. While not posing substantial technical climbing challenges on the standard route, Everest presents dangers such as altitude sickness, weather, and wind, as well as hazards from avalanches and the Khumbu Icefall. As of November 2022, 310 people have died on Everest. Over 200 bodies remain on the mountain and have not been removed due to the dangerous conditions. The first recorded efforts to reach Everest's summit were made by British mountaineers. As Nepal did not allow foreigners to enter the country at the time, the British made several attempts on the north ridge route from the Tibetan side. After the first reconnaissance expedition by the British in 1921 reached 7,000 m (22,970 ft) on the North Col, the 1922 expedition pushed the north ridge route up to 8,320 m (27,300 ft), marking the first time a human had climbed above 8,000 m (26,247 ft). The 1924 expedition resulted in one of the greatest mysteries on Everest to this day: George Mallory and Andrew Irvine made a final summit attempt on 8 June but never returned, sparking debate as to whether they were the first to reach the top. Tenzing Norgay and Edmund Hillary made the first documented ascent of Everest in 1953, using the southeast ridge route. Norgay had reached 8,595 m (28,199 ft) the previous year as a member of the 1952 Swiss expedition. The Chinese mountaineering team of Wang Fuzhou, Gonpo, and Qu Yinhua made the first reported ascent of the peak from the north ridge on 25 May 1960. \nQuestion: How high did they climb in 1922? According to the text, the 1922 expedition reached 8,"""

input_ids = tokenizer(prompt, return_tensors="pt", add_special_tokens=True).input_ids.to(model.device)
input_embeds = model.get_input_embeddings()(input_ids)
```

--------------------------------

### Clean Tokens and Generate Heatmap (Matplotlib/Pandas)

Source: https://lxt.readthedocs.io/en/latest/quickstart

This snippet processes the input token IDs into human-readable tokens, cleans them, and then generates a heatmap visualization of the computed relevance scores against these tokens. It utilizes a `pdf_heatmap` function, likely from a custom library, to create the visualization.

```python
tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
tokens = clean_tokens(tokens)
pdf_heatmap(tokens, relevance, path='heatmap_contrastive.pdf', backend='pdflatex')
```

--------------------------------

### Optimize LLaMA Model Memory Usage

Source: https://lxt.readthedocs.io/en/latest/quickstart

Disables gradients for model parameters to save memory and optionally enables gradient checkpointing. The model is set to training mode to allow gradient checkpointing, with LXT ensuring dropout rates are zeroed.

```python
# Deactivate gradients on parameters
for param in model.parameters():
    param.requires_grad = False

# Optionally enable gradient checkpointing (2x forward pass)
model.train()
model.gradient_checkpointing_enable()

```

--------------------------------

### Experiment with LRP Gamma Hyperparameters

Source: https://lxt.readthedocs.io/en/latest/quickstart

Iterates through different combinations of gamma hyperparameters for Conv2d and Linear layers using `itertools.product`. For each combination, it defines a `LayerMapComposite` with the specified gamma values, registers it with the model, performs a forward and backward pass, and calculates the LRP heatmap.

```python
# Store the generated heatmaps
heatmaps = []

# Experiment with different gamma values for Conv2d and Linear layers
# Gamma is a hyperparameter in LRP that controls how much positive vs. negative
# contributions are considered in the explanation
for conv_gamma, lin_gamma in itertools.product([0.1, 0.25, 100], [0, 0.01, 0.05, 0.1, 1]):
    input_tensor.grad = None  # Reset gradients
    print("Gamma Conv2d:", conv_gamma, "Gamma Linear:", lin_gamma)

    # Define rules for the Conv2d and Linear layers using 'zennit'
    # LayerMapComposite maps specific layer types to specific LRP rule implementations
    zennit_comp = LayerMapComposite([
        (torch.nn.Conv2d, z_rules.Gamma(conv_gamma)),
        (torch.nn.Linear, z_rules.Gamma(lin_gamma)),
    ])

    # Register the composite rules with the model
    zennit_comp.register(model)

    # Forward pass with gradient tracking enabled
    y = model(input_tensor.requires_grad_())

    # Get the top 5 predictions
    _, top5_classes = torch.topk(y, 5, dim=1)
    top5_classes = top5_classes.squeeze(0).tolist()

    # Get the class labels
    labels = weights.meta["categories"]
    top5_labels = [labels[class_idx] for class_idx in top5_classes]

    # Print the top 5 predictions and their labels
    for i, class_idx in enumerate(top5_classes):
        print(f'Top {i+1} predicted class: {class_idx}, label: {top5_labels[i]}')

    # Backward pass for the highest probability class
    # This initiates the LRP computation through the network
    y[0, top5_classes[0]].backward()

    # Remove the registered composite to prevent interference in future iterations
    zennit_comp.remove()

    # Calculate the relevance by computing Input*Gradient
    # This is the final step of LRP to get the pixel-wise explanation
    heatmap = (input_tensor * input_tensor.grad).sum(1)

    # Normalize relevance between [-1, 1] for plotting
    heatmap = heatmap / abs(heatmap).max()

    # Store the normalized heatmap
    heatmaps.append(heatmap[0].detach().cpu().numpy())

```

--------------------------------

### Save Generated Heatmaps

Source: https://lxt.readthedocs.io/en/latest/explicit_quickstart

Saves the collected heatmaps into a single image file named 'heatmap.png'. The `imgify` function from Zennit arranges the heatmaps into a grid with specified value limits (vmin=-1, vmax=1) and grid dimensions (3 rows, 5 columns).

```python
# save the heatmaps as a grid
imgify(heatmaps, vmin=-1, vmax=1, grid=(3, 5)).save('heatmap.png')

```

--------------------------------

### Compute and Normalize Relevance Scores (PyTorch)

Source: https://lxt.readthedocs.io/en/latest/quickstart

This code calculates the relevance scores for the input embeddings based on the gradients computed during the backward pass. The relevance scores are then normalized by their absolute maximum value to bring them into a standard range, making them easier to interpret.

```python
relevance = (input_embeds.grad * input_embeds).float().sum(-1).detach().cpu()[0]
relevance = relevance / relevance.abs().max()
```

--------------------------------

### Compute Input Embeddings and Gradients for Language Models (PyTorch)

Source: https://lxt.readthedocs.io/en/latest/quickstart

This snippet demonstrates how to compute input embeddings and their gradients with respect to the input for a language model using PyTorch. It handles the conversion to bfloat16 for gradient computation and activates gradient tracing for input embeddings. The output includes the computation of relevance scores.

```python
import torch
from transformers import AutoTokenizer

# Assuming tokenizer and model are already defined
# tokenizer = AutoTokenizer.from_pretrained(...)
# model = ...

# Define the prompt
prompt = """Context: Mount Everest attracts many climbers, including highly experienced mountaineers.\nThere are two main climbing routes, one approaching the summit from the southeast in Nepal (known as the standard route)\nand the other from the north in Tibet. While not posing substantial technical climbing challenges on the standard route,\nEverest presents dangers such as altitude sickness, weather, and wind, as well as hazards from avalanches and the Khumbu Icefall.\nAs of November 2022, 310 people have died on Everest. Over 200 bodies remain on the mountain and have not been removed\ndue to the dangerous conditions. The first recorded efforts to reach Everest's summit were made by British mountaineers.\nAs Nepal did not allow foreigners to enter the country at the time, the British made several attempts on the north ridge route\nfrom the Tibetan side. After the first reconnaissance expedition by the British in 1921 reached 7,000 m (22,970 ft) on the\nNorth Col, the 1922 expedition pushed the north ridge route up to 8,320 m (27,300 ft), marking the first time a human had\nclimbed above 8,000 m (26,247 ft). The 1924 expedition resulted in one of the greatest mysteries on Everest to this day:\nGeorge Mallory and Andrew Irvine made a final summit attempt on 8 June but never returned, sparking debate as to whether\nthey were the first to reach the top. Tenzing Norgay and Edmund Hillary made the first documented ascent of Everest in 1953,\nusing the southeast ridge route. Norgay had reached 8,595 m (28,199 ft) the previous year as a member of the 1952 Swiss expedition.\nThe Chinese mountaineering team of Wang Fuzhou, Gonpo, and Qu Yinhua made the first reported ascent of the peak from the\nnorth ridge on 25 May 1960.\nQuestion: How high did they climb in 1922? According to the text, the 1922 expedition reached 8,"""

# Get input embeddings
input_ids = tokenizer(prompt, return_tensors="pt", add_special_tokens=True).input_ids.to(model.device)
input_embeds = model.get_input_embeddings()(input_ids)

# Activate gradient tracing for the input embeddings
# Inference
output_logits = model(inputs_embeds=input_embeds.requires_grad_(), use_cache=False).logits

# Take the maximum logit at last token position.
max_logits, max_indices = torch.max(output_logits[0, -1, :], dim=-1)

# Backward pass (the relevance is initialized with the value of max_logits)
max_logits.backward()

# Obtain relevance.
relevance = (input_embeds.grad * input_embeds).float().sum(-1).detach().cpu()  # Cast to float32 for higher precision

```

--------------------------------

### Calculate BERT Relevance Scores and Generate Heatmap

Source: https://lxt.readthedocs.io/en/latest/quickstart

This snippet demonstrates how to calculate relevance scores for input tokens in a BERT model and visualize them as a heatmap. It involves tokenizing input, obtaining embeddings, calculating logits, and performing backpropagation to compute relevance. The output is a PDF heatmap visualizing token relevance.

```python
inputs = "After five years of research, scientists concluded that transformer models work because they has lots of parameters and math stuff."

input_ids = tokenizer(inputs, return_tensors="pt", add_special_tokens=True).input_ids.to("cuda")
inputs_embeds = model.bert.get_input_embeddings()(input_ids)

logits = model(inputs_embeds=inputs_embeds.requires_grad_(True)).logits

max_logits, max_indices = torch.max(logits, dim=-1)

out = model.config.id2label[max_indices.item()]
print("The label of the sequence is grammatically: ", out)

max_logits.backward()

relevance = (inputs_embeds * inputs_embeds.grad).float().sum(-1).detach().cpu()[0]

relevance = relevance / relevance.abs().max()

tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
tokens = clean_tokens(tokens)

pdf_heatmap(tokens, relevance, path="heatmap_bert.pdf", backend="xelatex")
```

--------------------------------

### Normalize and Visualize Relevance Scores in LaTeX (Python)

Source: https://lxt.readthedocs.io/en/latest/quickstart

This snippet shows how to normalize relevance scores to a [-1, 1] range and prepare tokens for LaTeX visualization. It then uses a helper function `pdf_heatmap` to generate a PDF heatmap of the relevance scores using the `xelatex` backend. This is useful for visualizing attention or importance scores in NLP tasks.

```python
# Normalize relevance between [-1, 1]
relevance = relevance / relevance.abs().max()

# Remove special characters that are not compatible wiht LaTeX
tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
tokens = clean_tokens(tokens)

# Save heatmap as PDF
pdf_heatmap(tokens, relevance[0], path='llama_heatmap.pdf', backend='xelatex')

```

--------------------------------

### Deactivate Gradients for Memory Saving (PyTorch)

Source: https://lxt.readthedocs.io/en/latest/quickstart

This snippet shows how to disable gradient calculation for model parameters in PyTorch to reduce memory consumption during inference or fine-tuning. It iterates through all parameters of a given model and sets their `requires_grad` attribute to `False`.

```python
for param in model.parameters():
    param.requires_grad = False
```

--------------------------------

### Using LXT Modules with Composite in PyTorch

Source: https://lxt.readthedocs.io/en/latest/on-the-fly

Illustrates an alternative way to use the Composite class by specifying LXT modules directly, such as LinearEpsilon and RMSNormIdentity, instead of LXT rules. This approach achieves a similar outcome of applying specific LXT functionalities to model layers.

```python
import lxt.modules
from lxt.core import Composite
import torch.nn as nn

# Assuming RootMeanSquareNorm is defined elsewhere
# class RootMeanSquareNorm(nn.Module):
#     pass

model = nn.Sequential(
    nn.Linear(10, 10),
    RootMeanSquareNorm(),
)

lrp = Composite({
        nn.Linear: lxt.modules.LinearEpsilon,
        RootMeanSquareNorm: lxt.modules.RMSNormIdentity,
    })

# To apply these modifications, you would typically call lrp.register(model)
# and to revert, lrp.remove()

```

--------------------------------

### Patching Hugging Face Models with LXT

Source: https://lxt.readthedocs.io/en/latest/extending

Demonstrates how to apply LXT's patching functionality to Hugging Face models for optimization. This involves patching attention functions, RMS-Norm layers, and MLP blocks to enforce uniform and identity rules.

```python
# 1. Patch attention functions
lxt.efficient.patches.patch_attention(eager_attention_forward, ALL_ATTENTION_FUNCTIONS)

# 2. Patch RMS-Norm forward pass
# (Assumes RMS-Norm class is accessible)
# Example: RMSNorm.forward = lxt.efficient.patches.patch_rms_norm_forward(RMSNorm.forward)

# 3. Patch MLP Block forward pass
# (Assumes MLPBlock class is accessible)
# Example: MLPBlock.forward = lxt.efficient.patches.patch_mlp_block_forward(MLPBlock.forward)

```

--------------------------------

### Automatic Operation Replacement with torch.fx and LXT

Source: https://lxt.readthedocs.io/en/latest/on-the-fly

This snippet demonstrates how to use torch.fx to automatically replace operations within a PyTorch model. It requires LXT and a dummy input to trace the model. The output is the traced model, which has operations modified according to the defined relevance propagation rules.

```python
import torch
import operator
import lxt

class SimpleModel(torch.nn.Module):

    def __init__(self):
        super().__init__()

        self.layer1 = torch.nn.Linear(10, 20, True)
        self.layer2 = torch.nn.Linear(10, 20, True)

    def forward(self, x):

        y1 = self.layer1(x)
        y2 = self.layer2(x)

        y1 = torch.nn.functional.softmax(y1, -1)

        return y1 + y2

model = SimpleModel()

lrp = Composite({
    nn.Linear: lxt.rules.EpsilonRule,
    operator.add: lxt.functional.add2,
    torch.nn.functional.softmax: lxt.functional.softmax,
})

x = torch.randn(1, 32, 10, requires_grad=True)
traced = lrp.register(model, dummy_inputs={'x': x}, verbose=True)

print(traced)
```

--------------------------------

### Applying LRP Rules with Composite in PyTorch

Source: https://lxt.readthedocs.io/en/latest/on-the-fly

Demonstrates how to use the Composite class to apply LRP rules (like EpsilonRule and IdentityRule) to specific module types within a PyTorch model. The Composite class wraps modules, and the register method applies these rules. The code shows how to instantiate Composite, register it with a model, and then use the modified model for computation and backpropagation.

```python
from lxt.core import Composite
import lxt.rules as rules
import torch.nn as nn

# Assuming RootMeanSquareNorm is defined elsewhere
# class RootMeanSquareNorm(nn.Module):
#     pass

model = nn.Sequential(
    nn.Linear(10, 10),
    RootMeanSquareNorm(),
)

lrp = Composite({
    nn.Linear: rules.EpsilonRule,
    RootMeanSquareNorm: rules.IdentityRule,
}, verbose=True)

# wrap modules in LXT rules and show the progress
lrp.register(model, verbose=True)

# print model to see the modification
print(model)

# Assuming x is defined and requires gradients
# x = torch.randn(1, 10, requires_grad=True)
# y = model(x.requires_grad_())
# y.max().backward()
# relevance = x.grad

```

--------------------------------

### Apply Uniform Rule with Scaled Dot-Product Attention (Python)

Source: https://lxt.readthedocs.io/en/latest/extending

This snippet demonstrates applying the uniform rule to inputs and outputs of scaled dot-product attention, including query and key matrices, and the final attention output. It utilizes `lxt.efficient.rules.divide_gradient` for this purpose.

```python
from lxt.efficient.rules import divide_gradient

# some attention function
def forward(self, key, value, query):
    # ....
    query = divide_gradient(query, 2) # Apply the uniform rule for query @ key multiplication
    key = divide_gradient(key, 2) # Apply the uniform rule for query @ key multiplication

    attn_output = torch.nn.functional.scaled_dot_product_attention(
        query,
        key,
        value,
        attn_mask=attention_mask,
        dropout_p=0.0, # No dropout if model is in train() mode
        is_causal=is_causal,
    )

    attn_output = divide_gradient(attn_output, 2) # Apply the uniform rule for softmax @ value multiplication
    # ....
```

--------------------------------

### Replace torch.nn.functional.linear with lxt.explicit.functional.linear_epsilon

Source: https://lxt.readthedocs.io/en/latest/drop-in-replacement

Demonstrates replacing `torch.nn.functional.linear` with `lxt.explicit.functional.linear_epsilon` to compute ε-LRP attributions for a linear function. This involves calling the LXT functional and then performing a backward pass to obtain relevance scores.

```python
import lxt.explicit.functional as lf

y = lf.linear_epsilon(x, W, b)

# initialize relevance with y itself
y.backward(y)

relevance = x.grad

# or for instance explain max output only
y = lf.linear_epsilon(x, W, b)
y.max().backward()

relevance = x.grad
```

--------------------------------

### Super-function for Epsilon-LRP using Vector-Jacobian Products

Source: https://lxt.readthedocs.io/en/latest/under-the-hood

Defines an abstract super-function to compute epsilon-LRP for arbitrary torch.nn.Module instances using PyTorch's vector-Jacobian products. This function wraps modules, avoiding the need to replace individual operations.

```python
def my_super_function(module, inputs, out_relevance, epsilon=1e-6):

    outputs = module(inputs)

    relevance_norm = out_relevance / (outputs + epsilon)

    # computes vector-jacobian product
    grads = torch.autograd.grad(outputs, inputs, relevance_norm)

    relevance = grads * inputs
    return relevance
```

--------------------------------

### Trace Relevance Flow in LLaMA Model using PyTorch Hooks

Source: https://lxt.readthedocs.io/en/latest/latent-feature-attribution-efficient

This code snippet demonstrates how to trace the relevance flow through a LLaMA model using PyTorch hooks. It involves setting up the model, tokenizer, applying forward hooks to capture activations and gradients, performing a forward and backward pass to compute relevance, and finally visualizing the relevance trace as a heatmap. Dependencies include PyTorch, Transformers, Matplotlib, and NumPy.

```python
import torch
from transformers import AutoTokenizer
from transformers.models.llama import modeling_llama
import matplotlib.pyplot as plt
import numpy as np

from lxt.efficient import monkey_patch
monkey_patch(modeling_llama, verbose=True)


def save_heatmap(values, tokens, figsize, title, save_path):
    fig, ax = plt.subplots(figsize=figsize)

    abs_max = abs(values).max()
    im = ax.imshow(values, cmap='bwr', vmin=-abs_max, vmax=abs_max)

    layers = np.arange(values.shape[-1])

    ax.set_xticks(np.arange(len(layers)))
    ax.set_yticks(np.arange(len(tokens)))

    ax.set_xticklabels(layers)
    ax.set_yticklabels(tokens)

    plt.title(title)
    plt.xlabel('Layers')
    plt.ylabel('Tokens')
    plt.colorbar(im)

    plt.show()
    plt.savefig(save_path, dpi=300, bbox_inches='tight')

def hook_hidden_activation(module, input, output):
    if isinstance(output, tuple):
        output = output[0]

    # save the activation and make sure the gradient is also saved in the .grad attribute after the backward pass
    module.output = output
    module.output.retain_grad() if module.output.requires_grad else None

model = modeling_llama.LlamaForCausalLM.from_pretrained('meta-llama/Llama-3.1-8B-Instruct', device_map='cuda', torch_dtype=torch.bfloat16)
tokenizer = AutoTokenizer.from_pretrained('meta-llama/Llama-3.1-8B-Instruct')

# optional gradient checkpointing to save memory (2x forward pass)
model.train()
model.gradient_checkpointing_enable()

# deactive gradients on parameters to save memory
for param in model.parameters():
    param.requires_grad = False

# apply hooks
for layer in model.model.layers:
    layer.register_forward_hook(hook_hidden_activation)

# forward & backard pass
prompt_response = f"I have 5 cats and 3 dogs. My cats love to play with my"
input_ids = tokenizer(prompt_response, return_tensors="pt", add_special_tokens=True).input_ids.to(model.device)
input_embeds = model.get_input_embeddings()(input_ids)

output_logits = model(inputs_embeds=input_embeds.requires_grad_(), use_cache=False).logits
max_logits, max_indices = torch.max(output_logits[:, -1, :], dim=-1)
max_logits.backward(max_logits)

print("Prediction:", tokenizer.convert_ids_to_tokens(max_indices))

# trace relevance through layers
relevance_trace = []
for layer in model.model.layers:
    relevance = (layer.output * layer.output.grad).float().sum(-1).detach().cpu()
    # normalize relevance at each layer between -1 and 1
    relevance = relevance / relevance.abs().max()
    relevance_trace.append(relevance)

relevance_trace = torch.cat(relevance_trace, dim=0)

tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
save_heatmap(relevance_trace.numpy().T, tokens, (20, 10), f"Latent Relevance Trace (Normalized)", f'latent_rel_trace.png')

```
Tokens:
Raw