Spaces:
Runtime error
Runtime error
File size: 3,461 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 |
"""
Runs statistic Cosine Similarity of Weights on all tensors of two given models, if they match in size.
Part of the "Unconstrained Setting" experiments (see StripedHyena experiments).
Relevant for hybrid models where only some parameters are shared.
To run: Use the HuggingFace Ids for the two models in Line 65-66.
Prints p-values between tensors that align in dimension.
"""
import torch
import scipy
from scipy.optimize import linear_sum_assignment as LAP
from transformers import AutoModelForCausalLM
from tracing.utils.utils import cossim, fisher
import warnings
warnings.filterwarnings("ignore")
def csw_sp_pair(base_model, ft_model, layer_name_base, layer_name_ft):
"""
Calculate Cosine Similarity of Weights between two specific layers.
Uses linear assignment to find optimal matching between neurons and
calculates Pearson correlation to quantify similarity.
Args:
base_model: First model to compare
ft_model: Second model to compare
layer_name_base: Name of the layer in the first model's state dict
layer_name_ft: Name of the layer in the second 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.pearsonr(matched.tolist(), orig.tolist())
return pvalue
def csw_models(base_model, ft_model):
"""
Perform comprehensive pairwise comparisons between all compatible layers of two models.
Tests all possible layer pairings between models that have compatible shapes,
useful for exploring model structure similarities without assuming corresponding positions.
Args:
base_model: First model to compare
ft_model: Second model to compare
Returns:
float: Aggregate p-value from Fisher's method combining all layer comparisons,
or 999 if no compatible layers were found
"""
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()):
# print(name1,name2)
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)
res = 0
if len(pvalues) == 0:
res = 999
else:
res = fisher(pvalues)
return res
def main():
model_1_id = "openai-community/gpt2"
model_2_id = "trl-internal-testing/dummy-GPT2-correct-vocab"
print(model_1_id, model_2_id)
model_1 = AutoModelForCausalLM.from_pretrained(model_1_id, torch_dtype=torch.bfloat16)
model_2 = AutoModelForCausalLM.from_pretrained(model_2_id, torch_dtype=torch.bfloat16)
print(csw_models(model_1, model_2))
if __name__ == "__main__":
main()
|