Spaces:
Runtime error
Runtime error
File size: 5,185 Bytes
de071e9 |
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 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 |
"""
Implementation of Cosine Similarity of Weights (CSW) test for comparing neural network models.
This module provides functions to test whether two models have similar weight matrices
using cosine similarity and statistical tests to quantify the similarity.
"""
import torch
from tracing.utils.utils import cossim, fisher
import scipy
import numpy as np
from scipy.stats import chi2
from scipy.optimize import linear_sum_assignment as LAP
def statistic(base_model, ft_model):
"""
Compute Cosine Similarity of Weights statistic between two models.
Args:
base_model: Base model to compare
ft_model: Fine-tuned or target model to compare against the base model
Returns:
tuple: (aggregate_p_value, p_values_per_layer) from the CSW test
"""
return csw_sp(base_model, ft_model)
def csw_sp_layer(base_model, ft_model, layer_name):
"""
Calculate Cosine Similarity of Weights for a specific layer.
Uses linear assignment to find optimal matching between neurons in the layer
and calculates Spearman correlation to quantify similarity.
Args:
base_model: Base model to compare
ft_model: Fine-tuned or target model to compare against the base model
layer_name: Name of the layer in the model's state dict to analyze
Returns:
float: p-value indicating the statistical similarity of weight matrices
"""
base_mat = base_model.state_dict()[layer_name]
ft_mat = ft_model.state_dict()[layer_name]
matched = LAP(cossim(base_mat.type(torch.float64), ft_mat.type(torch.float64)), maximize=True)
matched = matched[1]
orig = torch.arange(len(matched))
cor, pvalue = scipy.stats.spearmanr(matched.tolist(), orig.tolist())
return pvalue
def csw_sp(model1, model2):
"""
Apply CSW test across all MLP up-projection layers in the models.
Performs Fisher's method to combine p-values from individual layer tests
into an aggregate statistic.
Args:
model1: First model to compare
model2: Second model to compare
Returns:
tuple: (aggregate_p_value, list_of_p_values_per_layer)
"""
chi_squared = 0
num_layers = 0
p_values = []
for name1, name2 in zip(list(model1.state_dict().keys()), list(model2.state_dict().keys())):
if name1 != name2:
raise ValueError(f"Model parameter names do not match: {name1} != {name2}")
elif "mlp.up_proj" not in name1:
continue
pvalue = csw_sp_layer(model1, model2, name1)
if not np.isnan(pvalue):
chi_squared -= 2 * np.log(pvalue)
num_layers += 1
p_values.append(pvalue)
print(name1, pvalue)
aggregate_pvalue = chi2.sf(chi_squared, df=2 * num_layers)
return aggregate_pvalue, p_values
def csw_sp_pair(base_model, ft_model, layer_name_base, layer_name_ft):
"""
Calculate Cosine Similarity of Weights between two specific layers.
Similar to csw_sp_layer but allows comparing layers with different names.
Args:
base_model: Base model to compare
ft_model: Fine-tuned or target model to compare against the base model
layer_name_base: Name of the layer in the base model's state dict
layer_name_ft: Name of the layer in the fine-tuned model's state dict
Returns:
float: p-value indicating the statistical similarity of weight matrices
"""
base_mat = base_model.state_dict()[layer_name_base]
ft_mat = ft_model.state_dict()[layer_name_ft]
matched = LAP(cossim(base_mat.type(torch.float64), ft_mat.type(torch.float64)), maximize=True)
matched = matched[1]
orig = torch.arange(len(matched))
cor, pvalue = scipy.stats.spearmanr(matched.tolist(), orig.tolist())
return pvalue
def statistic_all(base_model, ft_model):
"""
Compute comprehensive pairwise comparisons between all compatible layers.
Tests every possible layer pairing between models that have compatible shapes,
useful for exploring model structure similarities without assumptions.
Args:
base_model: Base model to compare
ft_model: Fine-tuned or target model to compare against the base model
Returns:
float: Aggregate p-value from Fisher's method combining all layer comparisons
"""
base_model.to("cpu")
ft_model.to("cpu")
weights_base = base_model.state_dict()
weights_ft = ft_model.state_dict()
shapes_base = {}
shapes_ft = {}
for name1 in list(weights_base.keys()):
shapes_base[name1] = weights_base[name1].shape
for name2 in list(weights_ft.keys()):
shapes_ft[name2] = weights_ft[name2].shape
pvalues = []
for name1 in list(weights_base.keys()):
for name2 in list(weights_ft.keys()):
if shapes_base[name1] == shapes_ft[name2] and len(shapes_base[name1]) != 1:
pval = csw_sp_pair(base_model, ft_model, name1, name2)
print(name1, name2, pval)
pvalues.append(pval)
print(pvalues)
res = 0
if len(pvalues) == 0:
res = 999
else:
res = fisher(pvalues)
print(res)
return res
|