entropy-profiler 0.2.0

Extract, analyze, and visualize entropy profiles from transformer models using the logit-lens technique.

entropy-profiler

Extract, analyse, and visualize entropy profiles from transformer models using the logit-lens technique.

entropy-profiler computes per-layer Shannon or Rényi entropy by passing hidden states through the model's own unembedding head (layer norm + lm_head). It works on any HuggingFace CausalLM without architecture-specific hooks.

from entropy_profiler import EntropyProfiler, plot_profile
import torch

profiler = EntropyProfiler("gpt2", dtype=torch.float32)
profile = profiler.profile_text("The meaning of life is", max_new_tokens=32)
plot_profile(profile)

---

Installation

From source (recommended for development)

git clone https://github.com/TODO/entropy-profiler
cd entropy-profiler

# Using uv (fast, handles venvs automatically)
uv sync                      # core dependencies
uv sync --extra notebook     # + Jupyter support
uv sync --extra dev          # + pytest, ruff

# Or using pip
pip install -e .
pip install -e ".[notebook]"
pip install -e ".[dev]"

From PyPI (once published)

pip install entropy-profiler

---

Quick Start

Profile a single prompt

from entropy_profiler import EntropyProfiler, plot_profile
import torch

profiler = EntropyProfiler("gpt2", dtype=torch.float32)
profile = profiler.profile_text("The capital of France is", max_new_tokens=32)

print(profile.entropy.shape)    # (n_tokens, n_layers)
print(profile.mean_profile())   # (n_layers,) tensor
plot_profile(profile)

Profile multiple prompts

from entropy_profiler import plot_aggregated

agg = profiler.profile_batch([
    "The stock market experienced significant",
    "In quantum mechanics, the wave function",
    "Modern neural networks learn by",
], max_new_tokens=24)

print(agg.to_matrix().shape)    # (3, n_layers)
plot_aggregated(agg)

Compare prompts with distances

from entropy_profiler import profile_distance

p1 = profiler.profile_text("Water boils at", max_new_tokens=24)
p2 = profiler.profile_text("Once upon a time", max_new_tokens=24)

result = profile_distance(p1, p2, metric="jsd")
print(f"JSD distance: {result.aggregate:.4f}")

Analyse layer dynamics

from entropy_profiler import LayerAnalyzer

profile, hidden_states = profiler.profile_text_with_states(
    "Hello world", max_new_tokens=32
)
analyzer = LayerAnalyzer(profiler, profile, hidden_states=hidden_states)

print(analyzer.layer_entropy())          # (n_layers,)
print(analyzer.information_velocity())   # (n_layers,)
print(analyzer.layer_mi(method="cka"))   # (n_layers,)

---

Core Concepts

Logit-Lens Decoding

At each transformer layer, the hidden state is projected through the model's final layer norm and language model head to produce a vocabulary distribution. The entropy of this distribution measures how "decided" the model is at that layer — low entropy means a peaked distribution (confident prediction), high entropy means a flat distribution (uncertain).

Entropy Profiles

An entropy profile is a matrix of shape (n_tokens, n_layers) where each entry is the entropy of the vocabulary distribution at that token position and layer depth. The mean profile (n_layers,) averages across tokens to give a single curve showing how entropy evolves through the network.

Why Rényi Entropy?

Shannon entropy (alpha=1) is the default, but Rényi entropy at other orders provides complementary views:

• alpha < 1 — emphasises rare events (tail sensitivity)
• alpha = 1 — Shannon entropy (standard)
• alpha = 2 — collision entropy (sensitive to mode)
• alpha > 2 — increasingly dominated by the most probable token

---

API Reference

Core Module (entropy_profiler.profiler)

Symbol	Description
EntropyProfiler(model, dtype, alpha, layer_stride)	Main class. Loads model, runs generation, computes entropy.
EntropyProfile	Dataclass: entropy, token_ids, layer_indices, alpha, model_name, metadata.
AggregatedProfile	Collection of profiles with mean_profile() and to_matrix().
shannon_entropy(probs)	H(p) = -sum(p log p) on the last dimension.
renyi_entropy(probs, alpha)	Rényi entropy of order α. Falls back to Shannon when α ≈ 1.

EntropyProfiler methods:

Method	Returns	Description
profile_text(prompt, max_new_tokens, ...)	EntropyProfile	Profile generated text.
profile_text_with_states(prompt, ...)	(EntropyProfile, Tensor)	Profile + raw hidden states.
profile_batch(prompts, ...)	AggregatedProfile	Profile multiple prompts.
unload()	None	Free model memory.

EntropyProfile attributes and methods:

Member	Type	Description
entropy	Tensor (n_tokens, n_layers)	Per-token, per-layer entropy.
token_ids	Tensor (n_tokens,)	Generated token IDs.
n_layers	int	Number of profiled layers.
n_tokens	int	Number of profiled tokens.
mean_profile()	Tensor (n_layers,)	Mean entropy at each layer.
to_numpy()	ndarray	Convert to NumPy (float32).

Distances (entropy_profiler.distances)

Function	Type	Description
profile_distance(p1, p2, metric, aggregation)	DistanceResult	Unified entry point.
pairwise_distances(profiles, metric)	ndarray (N, N)	Symmetric distance matrix.
jsd_layer(p1, p2, n_bins)	ndarray (n_layers,)	Per-layer Jensen-Shannon divergence.
wasserstein_layer(p1, p2)	ndarray (n_layers,)	Per-layer Wasserstein-1 distance.
fisher_rao_distance(p1, p2)	float	Geodesic on probability simplex.
srvf_distance(p1, p2)	float	Elastic SRVF curve distance.

Available metrics for profile_distance: "jsd", "wasserstein", "fisher_rao", "srvf".

Aggregation methods: "mean", "max", "sum" (for layer-wise metrics).

Layer Analysis (entropy_profiler.analysis)

Symbol	Description
LayerAnalyzer(profiler, profile, hidden_states)	Per-layer metric computation.

Additional functions available via from entropy_profiler.analysis import ...: compare_models, plot_layer_importance, plot_information_plane, plot_velocity_entropy.

LayerAnalyzer methods:

Method	Returns	Description
layer_entropy()	ndarray (n_layers,)	Mean Shannon entropy per layer.
information_velocity()	ndarray (n_layers,)	Wasserstein between consecutive layers.
distance_to_output()	ndarray (n_layers,)	Fisher-Rao distance to final layer.
jsd_to_output(n_bins)	ndarray (n_layers,)	JSD from each layer to final.
layer_mi(method)	ndarray (n_layers,)	MI with final layer (Rényi or CKA).
layer_importance()	dict	All four non-MI metrics.

Visualization (entropy_profiler.viz)

Function	Description
plot_profile(profile, ax, ...)	Line plot with ±1 std fill.
plot_profiles(profiles, labels, ...)	Overlay multiple profiles.
plot_heatmap(profile, ax, ...)	Token × layer entropy heatmap.
plot_aggregated(agg, ax, ...)	Aggregated mean ± std curve.
plot_cluster(profiles, labels, method, feature, metric, ...)	2D scatter via t-SNE/UMAP/PCA.

Estimators (entropy_profiler.estimators)

Symbol	Description
MatrixRenyiMI(alpha, device)	Matrix-based Rényi MI via Gram matrices. Used by LayerAnalyzer.layer_mi().

---

Supported Models

Any HuggingFace AutoModelForCausalLM is supported. The profiler automatically detects the unembedding architecture:

Model Family	Layer Norm Path	Status
GPT-2	transformer.ln_f	Tested
LLaMA / LLaMA 2 / LLaMA 3	model.norm	Tested
Mistral	model.norm	Tested
Gemma / Gemma 2	model.norm	Tested
Qwen / Qwen 2	model.norm	Tested
OPT	model.norm (fallback)	Tested

To add a new architecture, add a resolution pattern to _get_unembedding() in entropy_profiler/profiler.py.

Tips for large models

# Use float16 for 7B+ models to fit in GPU memory
profiler = EntropyProfiler("meta-llama/Llama-2-7b-hf", dtype=torch.float16)

# Profile every other layer to reduce computation
profiler = EntropyProfiler("gpt2", layer_stride=2)

# Use context manager to auto-unload
with EntropyProfiler("gpt2") as profiler:
    profile = profiler.profile_text("Hello world")

---

Design Decisions

No hooks. HuggingFace's output_hidden_states=True returns all hidden states without hook infrastructure. This works across all CausalLM architectures with zero architecture-specific code.

Logit-lens, not probing. The unembedding head is the model's own decoder. No training of linear probes is needed — the entropy values are directly interpretable as "how peaked is the vocabulary distribution at this layer."

Float32 entropy. Entropy is always computed in float32 regardless of model dtype, avoiding numerical issues with half-precision softmax.

Dataclass outputs. EntropyProfile and DistanceResult are plain dataclasses — easy to inspect, serialize, and compose.

---

Notebooks

Notebook	Description
exploration.ipynb	Multi-model exploration: entropy curves, heatmaps, velocities, distances, clustering. Requires GPU and gated-model access.
api_tour.ipynb	Complete API tour exercising every public function with GPT-2.

uv sync --extra notebook
uv run jupyter notebook notebooks/

---

Development

# Install dev dependencies
uv sync --extra dev

# Lint
uv run ruff check .
uv run ruff check --fix .

# Test
uv run pytest

# Run a script without activating venv
uv run python your_script.py

---

License

MIT