File size: 3,524 Bytes
fed1832 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
#!/usr/bin/env python
import os
import torch
import numpy as np
import matplotlib.pyplot as plt
torch.set_printoptions(profile="full")
FILTER_RATE = 0.95
TOP_RATE = 0.5
ACTIVATION_BAR_RATIO = 0.95
langs = ["en", "eu"]
base_path = "new_activations"
n, over_zero = [], []
model_name = None
checkpoint = None
for lang in langs:
# File path
path = os.path.join(base_path, f"activation.{lang}.train.l2-7b-eu.pt")
data = torch.load(path)
n.append(data["n"])
over_zero.append(data["over_zero"])
# Extract model_name and checkpoint only once
if model_name is None:
model_name = os.path.basename(os.path.dirname(path)) # folder name
filename = os.path.basename(path)
parts = filename.split('.')
checkpoint = parts[-1] # 'qwen-checkpoint-1200'
# Convert to tensors
n = torch.Tensor(n) # (lang_num)
over_zero = torch.stack(over_zero, dim=-1) # (layer_num, neuron_num, lang_num)
num_layers, intermediate_size, lang_num = over_zero.size()
# 1. Activation probability
activation_probs = over_zero / n # broadcast
# 2. Normalized activation probability
normed_activation_probs = activation_probs / activation_probs.sum(dim=-1, keepdim=True)
# 3. LAPE (entropy)
log_prop = torch.where(normed_activation_probs > 0,
normed_activation_probs.log(),
torch.zeros_like(normed_activation_probs))
entropy = -(normed_activation_probs * log_prop).sum(dim=-1)
# 4. Filter neurons using 95th percentile
flat_probs = activation_probs.flatten()
thresh = flat_probs.kthvalue(int(flat_probs.numel() * FILTER_RATE)).values
valid_mask = (activation_probs > thresh).any(dim=-1) # [layers, neurons]
entropy[~valid_mask] = float("inf")
# 5. Select top-k neurons with lowest entropy
flat_entropy = entropy.flatten()
topk = int(flat_entropy.numel() * TOP_RATE)
_, idx = flat_entropy.topk(topk, largest=False)
layer_idx = idx // intermediate_size
neuron_idx = idx % intermediate_size
# 6. Group by languages
selection_props = activation_probs[layer_idx, neuron_idx] # [topk, lang_num]
bar = flat_probs.kthvalue(int(flat_probs.numel() * ACTIVATION_BAR_RATIO)).values
lang_mask = (selection_props > bar).T # [lang_num, topk]
final_mask = {}
for i, lang in enumerate(langs):
neuron_ids = torch.where(lang_mask[i])[0]
layer_lists = [[] for _ in range(num_layers)]
for nid in neuron_ids.tolist():
l = layer_idx[nid].item()
h = neuron_idx[nid].item()
layer_lists[l].append(h)
final_mask[lang] = [torch.tensor(lst, dtype=torch.long) for lst in layer_lists]
# =========================
# Plot number of neurons per layer (bar chart)
# =========================
plt.figure(figsize=(12, 6))
x = np.arange(num_layers)
width = 0.35
bars_list = []
for i, lang_key in enumerate(langs):
counts = [len(layer) for layer in final_mask[lang_key]]
bars = plt.bar(x + i * width, counts, width=width, label=lang_key)
bars_list.append(bars)
# Thêm số lên mỗi bar
for bar in bars:
height = bar.get_height()
plt.text(bar.get_x() + bar.get_width()/2.0, height, f'{int(height)}',
ha='center', va='bottom', fontsize=9)
plt.xlabel("Layer Index")
plt.ylabel("Number of Neurons")
plt.title(f"Number of Language-Specific Neurons per Layer\nModel: {model_name}, Checkpoint: {checkpoint}")
plt.xticks(x + width / 2, x)
plt.legend()
plt.grid(alpha=0.3, axis='y')
plt.tight_layout()
plt.savefig(f"{model_name}_{checkpoint}_neurons_bar.png", dpi=300)
plt.close()
|