File size: 1,854 Bytes
a96891a | 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 | import torch
import copy
import numpy as np
from scipy.stats import pearsonr
from t_cube import evaluate_model
def evaluate_slerp(clip_pt, sd_pt, sd_ft, dataloader, args, alpha=0.5):
"""
SLERP (spherical linear interpolation) between pretrained (pt) and fine-tuned (ft) weights.
alpha=0 -> pt only; alpha=1 -> ft only.
"""
model = copy.deepcopy(clip_pt)
merged_sd = {}
# flatten-per-key SLERP
for k in sd_pt.keys():
w1 = sd_pt[k].flatten().float()
w2 = sd_ft[k].flatten().float()
# cosine similarity
cos_val = torch.dot(w1, w2) / (w1.norm() * w2.norm() + 1e-8)
omega = torch.acos(torch.clamp(cos_val, -1+1e-6, 1-1e-6))
sin_omega = torch.sin(omega)
if sin_omega < 1e-6:
w_interp = (1-alpha)*w1 + alpha*w2
else:
w_interp = (torch.sin((1-alpha)*omega)/sin_omega)*w1 + \
(torch.sin(alpha*omega)/sin_omega)*w2
merged_sd[k] = w_interp.view_as(sd_pt[k])
model.load_state_dict(merged_sd)
return evaluate_model(model, dataloader, args)
def evaluate_m3(clip_pt, sd_pt, sd_ft, dataloader, args):
"""
M^3 (Mixup Model Merge): sample lambda ~ Uniform(0,1) and do linear interpolation.
"""
model = copy.deepcopy(clip_pt)
lam = np.random.rand()
merged_sd = {k: lam * sd_ft[k] + (1 - lam) * sd_pt[k]
for k in sd_pt.keys()}
model.load_state_dict(merged_sd)
return evaluate_model(model, dataloader, args)
def evaluate_task_arithmetic(clip_pt, sd_pt, sd_ft, dataloader, args):
"""
Task Arithmetic: extrapolate along the ft−pt vector, i.e. 2*ft – pt.
"""
model = copy.deepcopy(clip_pt)
merged_sd = {k: 2 * sd_ft[k] - sd_pt[k] for k in sd_pt.keys()}
model.load_state_dict(merged_sd)
return evaluate_model(model, dataloader, args)
|