llm_mri 0.1.10

Package to visualize LLM's Neural Networks activation regions

LLM-MRI: a brain scanner for LLMs

This repository contains the implementation from the paper LLM-MRI Python module: a brain scanner for LLMs

As the everyday use of large language models (LLMs) expands, so does the necessity of understanding how these models achieve their designated outputs. While many approaches focus on the interpretability of LLMs through visualizing different attention mechanisms and methods that explain the model's architecture, LLM-MRI focuses on the activations of the feed-forward layers in a transformer-based LLM.

By adopting this approach, the library examines the neuron activations produced by the model for each distinct label. Through a series of steps, such as dimensionality reduction and representing each layer as a grid, the tool provides various visualization methods for the activation patterns in the feed-forward layers. Accordingly, the objective of this library is to contribute to LLM interpretability research, enabling users to explore visualization methods, such as heatmaps and graph representations of the hidden layers' activations in transformer-based LLMs.

This model allows users to explore questions such as:

• How do different categories of text in the corpus activate different neural regions?
• What are the differences between the properties of graphs formed by activations from two distinct categories?
• Are there regions of activation in the model more related to specific aspects of a category?

We encourage you to not only use this toolkit but also to extend it as you see fit.

Index

• Online Example
• Installation
• Usage
• Functions
  • Activation Extraction
  • Heatmap Representation of Activations
  • Graph Representation of Activations
  • Composed Graph Visualization

Online Example

The link below runs an online example of our library, in the Jupyter platform running over the Binder server:

Instalation

To see LLM-MRI in action on your own data:

pip install llm_mri

Usage

Firstly, the user needs to import the LLM-MRI and matplotlib.pyplot packages:

from llm_mri import LLM_MRI
import matplotlib.pyplot as plt

The user also needs to specify the Hugging Face Dataset that will be used to process the model's activations. There are two ways to do this:

• Load the Dataset from Hugging Face Hub:

  dataset_url = "https://huggingface.co/datasets/dataset_link"
  dataset = load_dataset("csv", data_files=dataset_url)
  

• If you already have the dataset loaded on your machine, you can use the load_from_disk function:

  dataset = load_from_disk(dataset_path) # Specify the Dataset's path
  

Make sure that the selected dataset is a HuggingFace Dataset, and contains the columns "text" and "label", the last one being "ClassLabel" type. For more instructions on how to make this conversion, check out some of the examples on the GitHub documentation.

Next, the user selects the model to be used as a string:

model_ckpt = "distilbert/distilbert-base-multilingual-cased"

Then, the user instantiates LLM-MRI, to apply the methods defined on Functions:

llm_mri = LLM_MRI(model=model_ckpt, device="cpu", dataset=dataset)

For now, we recommend to use "cpu" as device. Further tests are going to be executed to ensure full "gpu" compatibility.

Functions

The library's functionality is divided into the following sections:

Activation Extraction:

As the user inputs the model and corpus to be analyzed, the dimensionality of the model's hidden layers is reduced, enabling visualization as an NxN grid.

llm_mri.process_activation_areas(map_dimension)

Heatmap representation of activations:

This includes the get_layer_image function, which transforms the NxN grid for a selected layer into a heatmap. In this heatmap, each cell represents the number of activations that different regions on a determined layer received for the provided corpus. Additionally, users can visualize activations for a specific label.

fig = llm_mri.get_layer_image(layer, category)

Graph Representation of Activations:

Using the get_graph function, the module connects regions from neighboring layers based on co-activations to form a graph representing the entire network. The graph's edges can also be colored according to different labels, allowing the user to identify the specific category that activated each neighboring node.

colormap: The default used colormap is the 'coolwarm'. More can be found on matplotlib.colors. We recommend the use of a 'Diverging' colormap for better visualization. fix_node_positions: 'True' keeps the nodes and edges at the same positions, independently of the categories. This could be useful for comparing activations between distinct categories. Setting to 'False' does not allow this comparison, although the graph will be more easily visualized.

graph = llm_mri.get_graph(category)
graph_image = llm_mri.get_graph_image(graph, colormap, fix_node_positions)

Composed Graph Visualization:

The user is also able to obtain a composed visualization of two different categories using the get_composed_graph function. By setting a category, each edge is colored based on the designated label, so the user is able to see which document label activated each region. Additionally, the user can select a colormap, where node colors reflect the label that most strongly activated them. Nodes colored white indicate equal activation by both categories, as white represents the midpoint of the color spectrum.

g_composed = llm_mri.get_composed_graph("true", "fake")
g_composed_img = llm_mri.get_graph_image(g_composed)

Reduced dimensionality representation of the documents

It is also possible to analyze where documents are disposed on a 2D space (after the dimensionality reduction) for a certain layer, divided by category. This could be essentialy more useful to visualize the impact of the DR algorithm on the activations disposal on the 2D space.

For comparison purpouses, this method differs from the heatmap representation of activations only by the cut of the 2D representation into a grid. With this approach, it is also posible to compare all documents on a single graph.

fig_scatter = llm_mri.get_original_map(layer=6)