3️⃣ White-box Methods

def extract_attention_matrix(
    text: str,
    model,
    tokenizer,
    layer: int,
    head: int,
) -> tuple[Float[np.ndarray, "seq seq"], list[str]]:
    """
    Extract attention matrix from a specific layer and head using hooks.
    Args:
        text: Input text to analyze
        model: HuggingFace model (already loaded)
        tokenizer: Corresponding tokenizer
        layer: Which layer to extract from (0-indexed)
        head: Which attention head to extract from (0-indexed)
    Returns:
        attention_matrix: Shape (seq_len, seq_len) attention weights for the specified head
        tokens: List of token strings for visualization
    """
# Tokenize input
    inputs = tokenizer(text, return_tensors="pt")
    if model.device.type == "cuda":
        inputs = {k: v.cuda() for k, v in inputs.items()}
# Forward pass - hook will capture attention
    with torch.no_grad():
        output = model(**inputs, use_cache=False, output_attentions=True)
# captured_attention[0] has shape (batch_size, num_heads, seq_len, seq_len)
    attn_tensor = output.attentions[layer][0, head]  # Select batch 0, specific head
# Convert to numpy
    attention_matrix = attn_tensor.cpu().numpy().astype(np.float32)
# Get token strings
    token_ids = inputs["input_ids"][0].tolist()
    tokens = tokenizer.convert_ids_to_tokens(token_ids)
return attention_matrix, tokens
# Load model and tokenizer (with attention output enabled - important!)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME_1B)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME_1B, device_map="auto")  # torch_dtype=torch.bfloat16
model.config._attn_implementation = "eager"
model.config.output_attentions = True
model.eval()
# Test with a short example
test_text = "The cat sat on the mat. It was sleeping."
# Extract attention from middle layer
n_layers = len(model.model.layers)
layer_to_check = n_layers // 2
head_to_check = 8  # head 9 is better for the 1.5b model
print(f"\nExtracting attention from head L{layer_to_check}H{head_to_check}")
attention_matrix, tokens = extract_attention_matrix(
    test_text, model, tokenizer, layer=layer_to_check, head=head_to_check
)
assert attention_matrix.shape[0] == len(tokens), "Shape mismatch"
assert attention_matrix.shape[0] == attention_matrix.shape[1], "Not square"
assert np.allclose(attention_matrix.sum(axis=1), 1.0, atol=1e-5), "Rows don't sum to 1"
# Check causal structure (upper triangle should be mostly zeros)
upper_triangle_sum = np.triu(attention_matrix, k=1).sum()
assert upper_triangle_sum < 1e-5, "Upper triangle has non-zero values, not causal attention"
# Visualize
display(
    cv.attention.attention_heads(
        attention=attention_matrix[None],
        tokens=tokens,
    )
)

def get_vertical_scores_simple(
    avg_attention_matrix: Float[np.ndarray, "n_sentences n_sentences"],
    proximity_ignore: int = 4,
) -> Float[np.ndarray, " n_sentences"]:
    """
    Compute basic vertical attention scores for each sentence (no normalization).
    Args:
        avg_attention_matrix: Shape (n_sentences, n_sentences), where entry (i, j)
            is the average attention from sentence i to sentence j
        proximity_ignore: Ignore this many nearby sentences (to avoid trivial patterns)
    Returns:
        Array of shape (n_sentences,) with vertical scores
    The vertical score for sentence j is the mean attention it receives from
    all sentences i where i > j + proximity_ignore.
    Adapted from thought-anchors: attn_funcs.py:get_vertical_scores
    """
    n = avg_attention_matrix.shape[0]
    mat = avg_attention_matrix.copy()
# Step 1: Clean matrix - set upper triangle to NaN (can't attend to future)
    mat[np.triu_indices_from(mat, k=1)] = np.nan
# Step 2: Remove proximity - set near-diagonal to NaN (ignore nearby sentences)
    mat[np.triu_indices_from(mat, k=-proximity_ignore + 1)] = np.nan
# Step 3: Compute vertical scores: mean attention received from future sentences
    vert_scores = []
    for j in range(n):
        # Extract "vertical line" - attention from all future sentences to sentence j
        future_attention = mat[j + proximity_ignore :, j]
if len(future_attention) == 0 or np.all(np.isnan(future_attention)):
            vert_scores.append(np.nan)
        else:
            vert_scores.append(np.nanmean(future_attention))
return np.array(vert_scores)

def get_vertical_scores(
    avg_attention_matrix: Float[np.ndarray, "n_sentences n_sentences"],
    proximity_ignore: int = 4,
    control_depth: bool = True,
    return_z_scores: bool = False,
) -> Float[np.ndarray, " n_sentences"]:
    """
    Compute vertical attention scores for each sentence with optional depth control and z-scoring.
    Args:
        avg_attention_matrix: Shape (n_sentences, n_sentences), where entry (i, j)
            is the average attention from sentence i to sentence j
        proximity_ignore: Ignore this many nearby sentences (to avoid trivial patterns)
        control_depth: Apply rank normalization to control for depth effects
        return_z_scores: If True, return z-score normalized values
    Returns:
        Array of shape (n_sentences,) with vertical scores
    The vertical score for sentence j is the mean attention it receives from
    all sentences i where i > j + proximity_ignore.
    Depth control (when enabled): Normalizes each row by ranking attention values
    and dividing by the count of valid (non-NaN) positions. This ensures fair
    comparison between early sentences (fewer tokens to attend to) and late
    sentences (many tokens to attend to).
    Adapted from thought-anchors: attn_funcs.py:get_vertical_scores
    """
    n = avg_attention_matrix.shape[0]
    mat = avg_attention_matrix.copy()
# Step 1: Clean matrix - set upper triangle to NaN (can't attend to future)
    mat[np.triu_indices_from(mat, k=1)] = np.nan
# Step 2: Remove proximity - set near-diagonal to NaN (ignore nearby sentences)
    mat[np.triu_indices_from(mat, k=-proximity_ignore + 1)] = np.nan
# Step 3: Depth control normalization (critical for fair comparison!)
    if control_depth:
        # Count non-NaN values per row (available positions to attend to)
        per_row = np.sum(~np.isnan(mat), axis=1)
# Rank-normalize each row: convert attention values to ranks,
        # then divide by number of valid positions
        # This puts all rows on the same scale regardless of depth
        mat = stats.rankdata(mat, axis=1, nan_policy="omit") / per_row[:, None]
# Step 4: Compute vertical scores: mean attention received from future sentences
    vert_scores = []
    for j in range(n):
        # Extract "vertical line" - attention from all future sentences to sentence j
        future_attention = mat[j + proximity_ignore :, j]
if len(future_attention) == 0 or np.all(np.isnan(future_attention)):
            vert_scores.append(np.nan)
        else:
            vert_scores.append(np.nanmean(future_attention))
vert_scores = np.array(vert_scores)
# Step 5: Optional z-score normalization
    if return_z_scores:
        vert_scores = (vert_scores - np.nanmean(vert_scores)) / np.nanstd(vert_scores)
return vert_scores

def compute_head_kurtosis(vertical_scores: Float[np.ndarray, " n_sentences"]) -> float:
    """
    Compute kurtosis of vertical attention scores.
    Higher kurtosis = more "spiky" attention = better receiver head candidate.
    """
    # Remove NaN values
    valid_scores = vertical_scores[~np.isnan(vertical_scores)]
    if len(valid_scores) < 4:
        return np.nan
# Compute Fisher's kurtosis (excess kurtosis)
    return stats.kurtosis(valid_scores, fisher=True, bias=True)

def find_receiver_heads(
    text: str,
    sentences: list[str],
    model,
    tokenizer,
    top_k: int = 10,
    proximity_ignore: int = 16,
) -> tuple[Int[np.ndarray, "top_k 2"], Float[np.ndarray, " top_k"]]:
    """
    Identify top-k receiver heads by kurtosis of vertical attention scores.
    Args:
        text: Full reasoning trace
        sentences: List of sentences in the trace
        model: Loaded model
        tokenizer: Corresponding tokenizer
        top_k: Number of top receiver heads to return
        proximity_ignore: Proximity parameter for vertical scores
    Returns:
        receiver_heads: Shape (top_k, 2) array of (layer, head) pairs
        kurtosis_scores: Shape (top_k,) array of kurtosis values for each head
    """
n_layers = len(model.model.layers)
    n_heads = model.config.num_attention_heads
    n_sentences = len(sentences)
print(f"Computing vertical scores for {n_layers} layers × {n_heads} heads = {n_layers * n_heads} heads...")
    print("This may take a few minutes...")
# Get sentence boundaries once (expensive operation)
    boundaries = utils.get_sentence_token_boundaries(text, sentences, tokenizer)
# Store vertical scores for all heads
    all_vert_scores = np.zeros((n_layers, n_heads, n_sentences))
# Extract attention and compute vertical scores for each head
    for layer in tqdm(range(n_layers), desc="Layers"):
        for head in range(n_heads):
            # Extract attention matrix
            attn_matrix, _ = extract_attention_matrix(text, model, tokenizer, layer, head)
# Average to sentence level
            avg_attn = utils.average_attention_by_sentences(attn_matrix, boundaries)
# Compute vertical scores with depth control
            vert_scores = get_vertical_scores(avg_attn, proximity_ignore=proximity_ignore, control_depth=True)
all_vert_scores[layer, head] = vert_scores
torch.cuda.empty_cache()
# Compute kurtosis for each head (over sentences dimension)
    kurtosis_matrix = np.zeros((n_layers, n_heads))
    for layer in range(n_layers):
        for head in range(n_heads):
            kurtosis_matrix[layer, head] = compute_head_kurtosis(all_vert_scores[layer, head])
# Find top-k heads with highest kurtosis
    flat_kurts = kurtosis_matrix.flatten()
    valid_indices = ~np.isnan(flat_kurts)
if valid_indices.sum() < top_k:
        print(f"Warning: Only {valid_indices.sum()} valid heads found")
        top_k = valid_indices.sum()
# Get indices of top-k highest kurtosis values
    valid_kurts = flat_kurts[valid_indices]
    valid_flat_indices = np.where(valid_indices)[0]
top_k_in_valid = np.argpartition(valid_kurts, -top_k)[-top_k:]
    top_k_in_valid = top_k_in_valid[np.argsort(-valid_kurts[top_k_in_valid])]  # Sort descending
top_flat_indices = valid_flat_indices[top_k_in_valid]
# Convert flat indices back to (layer, head) pairs
    receiver_heads = np.array(np.unravel_index(top_flat_indices, (n_layers, n_heads))).T
    receiver_kurts = flat_kurts[top_flat_indices]
return receiver_heads, receiver_kurts
# Load example CoT
text_full, sentences_full, _ = utils.get_whitebox_example_data(problem_data)
n_chunks = 74
sentences_subset = sentences_full[:n_chunks]
# Find where the n_chunks-th sentence ends in the original text to extract the correct substring
# This preserves the original text formatting and ensures tokenization consistency
end_char_pos = 0
for sent in sentences_subset:
    # Find this sentence in the text starting from where we left off
    sent_pos = text_full.find(sent, end_char_pos)
    if sent_pos == -1:
        # Try with stripped version
        sent_pos = text_full.find(sent.strip(), end_char_pos)
    if sent_pos != -1:
        end_char_pos = sent_pos + len(sent)
text_subset = text_full[:end_char_pos]
print(f"Analyzing first {len(sentences_subset)} sentences...")
print(f"Text length: {len(text_subset)} characters")
# Find top-20 receiver heads (paper recommends 16-32)
torch.cuda.empty_cache()
receiver_heads, receiver_kurts = find_receiver_heads(
    text_subset,
    sentences_subset,
    model,
    tokenizer,
    top_k=20,
    proximity_ignore=4,
)
print(f"\nTop-{len(receiver_heads)} Receiver Heads:")
for i, ((layer, head), kurt) in enumerate(zip(receiver_heads, receiver_kurts)):
    print(f"  {i + 1:2d}. Layer {layer:2d}, Head {head:2d} | Kurtosis: {kurt:.3f}")
# Visualize kurtosis distribution
# Recompute kurtosis matrix for visualization
n_layers = len(model.model.layers)
n_heads = model.config.num_attention_heads
kurtosis_viz = np.zeros((n_layers, n_heads))
for layer, head in receiver_heads:
    kurtosis_viz[layer, head] = receiver_kurts[receiver_heads.tolist().index([layer, head])]
fig = px.imshow(
    kurtosis_viz,
    color_continuous_scale="RdBu",
    color_continuous_midpoint=0.0,
    labels=dict(x="Head", y="Layer", color="Kurtosis"),
    title="Top Receiver Heads by Kurtosis",
    aspect="auto",
    width=800,
    height=600,
)
fig.show()
print("Receiver heads should attend 'spikily' to specific important sentences.")

def compute_receiver_head_scores(
    text: str,
    sentences: list[str],
    receiver_heads: Int[np.ndarray, "k 2"],
    model,
    tokenizer,
    proximity_ignore: int = 4,
) -> Float[np.ndarray, " n_sentences"]:
    """
    Compute final receiver head scores by averaging vertical scores from top-k heads.
    Args:
        text: Full reasoning trace
        sentences: List of sentences
        receiver_heads: Shape (k, 2) array of (layer, head) pairs
        model: Loaded model
        tokenizer: Corresponding tokenizer
        proximity_ignore: Proximity parameter
    Returns:
        receiver_scores: Shape (n_sentences,) - final importance scores
    """
boundaries = utils.get_sentence_token_boundaries(text, sentences, tokenizer)
all_vert_scores = []
print(f"Computing scores from {len(receiver_heads)} receiver heads...")
    for layer, head in receiver_heads:
        # Extract attention for this receiver head
        attn_matrix, _ = extract_attention_matrix(text, model, tokenizer, layer, head)
# Average to sentence level
        avg_attn = utils.average_attention_by_sentences(attn_matrix, boundaries)
# Compute vertical scores
        vert_scores = get_vertical_scores(avg_attn, proximity_ignore=proximity_ignore, control_depth=True)
all_vert_scores.append(vert_scores)
# Average across all receiver heads
    all_vert_scores = np.array(all_vert_scores)  # Shape: (k, n_sentences)
    receiver_scores = np.nanmean(all_vert_scores, axis=0)  # Average over heads
return receiver_scores
# Use the receiver heads we just found
print(f"Computing scores from {len(receiver_heads)} receiver heads...")
receiver_scores = compute_receiver_head_scores(
    text_subset,
    sentences_subset,
    receiver_heads,
    model,
    tokenizer,
    proximity_ignore=4,
)
print(f"\nReceiver head scores computed for {len(receiver_scores)} sentences")
print(f"Mean score: {np.nanmean(receiver_scores):.4f}")
print(f"Std score: {np.nanstd(receiver_scores):.4f}")
# Visualize
fig = go.Figure()
fig.add_trace(
    go.Bar(
        x=list(range(len(receiver_scores))),
        y=receiver_scores,
        marker_color="indianred",
        hovertemplate="<b>Sentence %{x}</b><br>Score: %{y:.4f}<br>Text: %{customdata}<extra></extra>",
        customdata=[s[:100] for s in sentences_subset],
    )
)
fig.add_hline(
    y=np.nanmean(receiver_scores),
    line_dash="dash",
    line_color="black",
    annotation_text="Mean",
)
fig.update_layout(
    title="Receiver Head Scores (Averaged Across Top-k Heads)",
    xaxis_title="Sentence Index",
    yaxis_title="Receiver Head Score",
    width=800,
    height=400,
)
fig.show()
print("\nTop-3 sentences by receiver head score:")
# Filter out NaN values before selecting top sentences
valid_mask = ~np.isnan(receiver_scores)
valid_indices = np.where(valid_mask)[0]
valid_scores = receiver_scores[valid_mask]
top_k = min(3, len(valid_scores))
top_in_valid = np.argsort(valid_scores)[-top_k:][::-1]
top_indices = valid_indices[top_in_valid]
for idx in top_indices:
    print(f"  Sentence {idx}: {receiver_scores[idx]:.4f}")
    print(f"    Text: {sentences_subset[idx][:100]}...")
print("Receiver head scores computed! These aggregate the attention from all high-kurtosis receiver heads.")

def rotate_half(x: Float[Tensor, "... d"]) -> Float[Tensor, " ... d"]:
    """Rotates half the hidden dims of the input (for RoPE)."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)
def apply_rotary_pos_emb(
    q: Float[Tensor, "batch q_heads seq head_dim"],
    k: Float[Tensor, "batch kv_heads seq head_dim"],
    cos: Float[Tensor, "batch seq head_dim"],
    sin: Float[Tensor, "batch seq head_dim"],
    position_ids: torch.Tensor | None = None,
    unsqueeze_dim: int = 1,
) -> tuple[Float[Tensor, "batch q_heads seq head_dim"], Float[Tensor, "batch kv_heads seq head_dim"]]:
    """Applies Rotary Position Embedding to query and key tensors."""
del position_ids  # unused
cos = cos.unsqueeze(unsqueeze_dim)  # Shape: (batch, 1, seq, head_dim)
    sin = sin.unsqueeze(unsqueeze_dim)  # Shape: (batch, 1, seq, head_dim)
    q_embed = (q  cos) + (rotate_half(q)  sin)
    k_embed = (k  cos) + (rotate_half(k)  sin)
    return q_embed, k_embed
tests.test_rotary_pos_emb(apply_rotary_pos_emb)

def compute_suppressed_attention(
    query_states: Float[Tensor, "batch heads seq head_dim"],
    key_states: Float[Tensor, "batch heads seq head_dim"],
    token_ranges: list[tuple[int, int]],
    head_dim: int,  # for scaling attention weights by sqrt of!
    query_len: int,  # in case token_ranges exceed query length
    attention_mask: Float[Tensor, "batch 1 seq seq"] | None = None,
    heads_mask: Int[Tensor, " heads"] | None = None,  # shape (num_heads,) with True for heads to suppress
) -> Float[Tensor, "batch heads seq seq"]:
    # Compute attention scores
    attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(head_dim)
# Apply suppression mask BEFORE softmax (key step!)
    for start, end in token_ranges:
        effective_start = min(start, query_len)
        effective_end = min(end, query_len)
        if effective_start < effective_end:
            mask_value = torch.finfo(attn_weights.dtype).min
            if heads_mask is None:
                # Suppress in all heads
                attn_weights[..., effective_start:effective_end] = mask_value
            else:
                # Suppress only in specific heads
                attn_weights[:, heads_mask, :, effective_start:effective_end] = mask_value
# Apply attention mask if provided
    if attention_mask is not None:
        attn_weights = attn_weights + attention_mask
# Softmax
    attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
    return attn_weights
tests.test_compute_suppressed_attention(compute_suppressed_attention)

def repeat_kv(
    hidden_states: Float[Tensor, "batch kv_heads seq head_dim"], n_rep: int
) -> Float[Tensor, "batch heads seq head_dim"]:
    """Expands key/value tensors for grouped-query attention."""
    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
def apply_qwen_attention_suppression(
    model,
    token_ranges: list[tuple[int, int]] | tuple[int, int],
    layer_to_heads: dict[int, list[int]] | None = None,
) -> dict[str, Any]:
    """
    Suppresses attention to specific token positions by replacing forward methods.
    Args:
        model: The model to apply suppression to
        token_ranges: Token range(s) to suppress - single tuple or list of tuples
        layer_to_heads: Dict mapping layer indices to lists of head indices to suppress (None = all heads)
    Returns:
        Dict with 'original_forwards' for restoration
    Adapted from thought-anchors: whitebox-analyses/pytorch_models/hooks.py
    """
    # Normalize token_ranges to list of tuples
    if isinstance(token_ranges, tuple):
        token_ranges = [token_ranges]
# Find rotary embedding module
    rotary_emb_module = None
    if hasattr(model, "model") and hasattr(model.model, "rotary_emb"):
        rotary_emb_module = model.model.rotary_emb
# Find attention modules to patch
    target_modules = []
    for name, module in model.named_modules():
        if name.startswith("model.layers") and name.endswith("self_attn"):
            try:
                layer_idx = int(name.split(".")[2])
                if layer_to_heads is None or layer_idx in layer_to_heads:
                    if all(hasattr(module, attr) for attr in ["config", "q_proj", "k_proj", "v_proj", "o_proj"]):
                        target_modules.append((name, module, layer_idx))
            except (IndexError, ValueError):
                continue
if not target_modules:
        print("Warning: No Qwen attention modules found to patch")
        return {"original_forwards": {}}
# Store original forward methods
    original_forwards = {}
# Create and apply masked forward functions
    for name, attn_module, layer_idx in target_modules:
        original_forwards[name] = attn_module.forward
        heads_mask = layer_to_heads[layer_idx] if layer_to_heads is not None else None
# Create masked forward function
        def create_masked_forward(orig_forward, layer_idx, rotary_ref, heads_mask):
            def masked_forward(
                self,
                hidden_states: torch.Tensor,
                attention_mask: torch.Tensor | None = None,
                position_ids: torch.LongTensor | None = None,
                past_key_value: tuple[torch.Tensor] | None = None,
                output_attentions: bool = False,
                use_cache: bool = False,
                cache_position: torch.LongTensor | None = None,
                **kwargs,
            ) -> tuple[torch.Tensor, torch.Tensor | None]:
                bsz, q_len, _ = hidden_states.size()
                config = self.config
                device = hidden_states.device
# Project to Q, K, V
                query_states = self.q_proj(hidden_states)
                key_states = self.k_proj(hidden_states)
                value_states = self.v_proj(hidden_states)
# Reshape for multi-head attention
                num_heads = config.num_attention_heads
                head_dim = config.hidden_size // num_heads
                num_key_value_heads = config.num_key_value_heads
                num_key_value_groups = num_heads // num_key_value_heads
query_states = query_states.view(bsz, q_len, num_heads, head_dim).transpose(1, 2)
                key_states = key_states.view(bsz, q_len, num_key_value_heads, head_dim).transpose(1, 2)
                value_states = value_states.view(bsz, q_len, num_key_value_heads, head_dim).transpose(1, 2)
# Apply RoPE
                if position_ids is None:
                    position_ids = torch.arange(0, q_len, dtype=torch.long, device=device).unsqueeze(0)
if rotary_ref is not None and callable(rotary_ref):
                    cos, sin = rotary_ref(value_states, position_ids=position_ids)
                    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
# Repeat K/V for grouped-query attention
                key_states = repeat_kv(key_states, num_key_value_groups)
                value_states = repeat_kv(value_states, num_key_value_groups)
attn_weights = compute_suppressed_attention(
                    query_states,
                    key_states,
                    token_ranges,
                    head_dim,
                    q_len,
                    attention_mask,
                    heads_mask,
                )
# Apply attention to values
                attn_output = torch.matmul(attn_weights, value_states)
# Reshape and project output
                attn_output = attn_output.transpose(1, 2).contiguous()
                attn_output = attn_output.reshape(bsz, q_len, config.hidden_size)
                attn_output = self.o_proj(attn_output)
return attn_output, attn_weights if output_attentions else None
return masked_forward
attn_module.forward = MethodType(
            create_masked_forward(attn_module.forward, layer_idx, rotary_emb_module, heads_mask), attn_module
        )
return {"original_forwards": original_forwards}
def remove_qwen_attention_suppression(model, suppression_info: dict[str, Any]):
    """Restores original forward methods after attention suppression."""
    original_forwards = suppression_info.get("original_forwards", {})
    if not original_forwards:
        return
for name, module in model.named_modules():
        if name in original_forwards:
            module.forward = original_forwards[name]

def compute_suppression_kl(
    original_logits: Float[np.ndarray, " vocab"],
    suppressed_logits: Float[np.ndarray, " vocab"],
    temperature: float = 1.0,
) -> float:
    """
    Compute KL divergence between original and suppressed output distributions.
    Args:
        original_logits: Logits from original forward pass
        suppressed_logits: Logits from suppressed forward pass
        temperature: Temperature for softmax
    Returns:
        KL divergence (in nats)
    """
def softmax(x, temp):
        x = x / temp
        exp_x = np.exp(x - np.max(x))
        return exp_x / exp_x.sum()
p = softmax(original_logits, temperature)
    q = softmax(suppressed_logits, temperature)
# Add small epsilon for numerical stability
    eps = 1e-10
    return np.sum(p * np.log((p + eps) / (q + eps)))