import math
import torch
from torch import Tensor
__all__ = [
"weighted_sum",
"weighted_subtraction",
"tensor_sum",
"add_difference",
"sum_twice",
"triple_sum",
"euclidean_add_difference",
"multiply_difference",
"top_k_tensor_sum",
"similarity_add_difference",
"distribution_crossover",
"ties_add_difference",
]
EPSILON = 1e-10
def weighted_sum(a: Tensor, b: Tensor, alpha: float, **kwargs) -> Tensor:
"""
Basic Merge:
alpha 0 returns Primary Model
alpha 1 returns Secondary Model
"""
return (1 - alpha) * a + alpha * b
def weighted_subtraction(a: Tensor, b: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
"""
The inverse of a Weighted Sum Merge
Returns Primary Model when alpha*beta = 0
High values of alpha*beta are likely to break the merged model
"""
if alpha == 1.0 and beta == 1.0:
beta -= EPSILON
return (a - alpha * beta * b) / (1 - alpha * beta)
def tensor_sum(a: Tensor, b: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
"""
Takes a slice of Secondary Model and pastes it into Primary Model
Alpha sets the width of the slice
Beta sets the start point of the slice
ie Alpha = 0.5 Beta = 0.25 is (ABBA) Alpha = 0.25 Beta = 0 is (BAAA)
"""
if alpha + beta <= 1:
tt = a.clone()
talphas = int(a.shape[0] * beta)
talphae = int(a.shape[0] * (alpha + beta))
tt[talphas:talphae] = b[talphas:talphae].clone()
else:
talphas = int(a.shape[0] * (alpha + beta - 1))
talphae = int(a.shape[0] * beta)
tt = b.clone()
tt[talphas:talphae] = a[talphas:talphae].clone()
return tt
def add_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, **kwargs) -> Tensor:
"""
Classic Add Difference Merge
"""
return a + alpha * (b - c)
def sum_twice(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
"""
Stacked Basic Merge:
Equivalent to Merging Primary and Secondary @ alpha
Then merging the result with Tertiary @ beta
"""
return (1 - beta) * ((1 - alpha) * a + alpha * b) + beta * c
def triple_sum(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
"""
Weights Secondary and Tertiary at alpha and beta respectively
Fills in the rest with Primary
Expect odd results if alpha + beta > 1 as Primary will be merged with a negative ratio
"""
return (1 - alpha - beta) * a + alpha * b + beta * c
def euclidean_add_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, **kwargs) -> Tensor:
"""
Subtract Primary and Secondary from Tertiary
Compare the remainders via Euclidean distance
Add to Tertiary
Note: Slow
"""
a_diff = a.float() - c.float()
b_diff = b.float() - c.float()
a_diff = torch.nan_to_num(a_diff / torch.linalg.norm(a_diff))
b_diff = torch.nan_to_num(b_diff / torch.linalg.norm(b_diff))
distance = (1 - alpha) * a_diff**2 + alpha * b_diff**2
distance = torch.sqrt(distance)
sum_diff = weighted_sum(a.float(), b.float(), alpha) - c.float()
distance = torch.copysign(distance, sum_diff)
target_norm = torch.linalg.norm(sum_diff)
return c + distance / torch.linalg.norm(distance) * target_norm
def multiply_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
"""
Similar to Add Difference but with geometric mean instead of arithmetic mean
"""
diff_a = torch.pow(torch.abs(a.float() - c), (1 - alpha))
diff_b = torch.pow(torch.abs(b.float() - c), alpha)
difference = torch.copysign(diff_a * diff_b, weighted_sum(a, b, beta) - c)
return c + difference.to(c.dtype)
def top_k_tensor_sum(a: Tensor, b: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
"""
Redistributes the largest weights of Secondary Model into Primary Model
"""
a_flat = torch.flatten(a)
a_dist = torch.msort(a_flat)
b_indices = torch.argsort(torch.flatten(b), stable=True)
redist_indices = torch.argsort(b_indices)
start_i, end_i, region_is_inverted = ratio_to_region(alpha, beta, torch.numel(a))
start_top_k = kth_abs_value(a_dist, start_i)
end_top_k = kth_abs_value(a_dist, end_i)
indices_mask = (start_top_k < torch.abs(a_dist)) & (torch.abs(a_dist) <= end_top_k)
if region_is_inverted:
indices_mask = ~indices_mask
indices_mask = torch.gather(indices_mask.float(), 0, redist_indices)
a_redist = torch.gather(a_dist, 0, redist_indices)
a_redist = (1 - indices_mask) * a_flat + indices_mask * a_redist
return a_redist.reshape_as(a)
def kth_abs_value(a: Tensor, k: int) -> Tensor:
if k <= 0:
return torch.tensor(-1, device=a.device)
else:
return torch.kthvalue(torch.abs(a.float()), k)[0]
def ratio_to_region(width: float, offset: float, n: int) -> tuple[int, int, bool]:
if width < 0:
offset += width
width = -width
width = min(width, 1)
if offset < 0:
offset = 1 + offset - int(offset)
offset = math.fmod(offset, 1.0)
if width + offset <= 1:
inverted = False
start = offset * n
end = (width + offset) * n
else:
inverted = True
start = (width + offset - 1) * n
end = offset * n
return round(start), round(end), inverted
def similarity_add_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
"""
Weighted Sum where A and B are similar and Add Difference where A and B are dissimilar
"""
threshold = torch.maximum(torch.abs(a), torch.abs(b))
similarity = ((a * b / threshold**2) + 1) / 2
similarity = torch.nan_to_num(similarity * beta, nan=beta)
ab_diff = a + alpha * (b - c)
ab_sum = (1 - alpha / 2) * a + (alpha / 2) * b
return (1 - similarity) * ab_diff + similarity * ab_sum
def distribution_crossover(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs):
"""
From the creator:
It's Primary high-passed + Secondary low-passed. Takes the fourrier transform of the weights of
Primary and Secondary when ordered with respect to Tertiary. Split the frequency domain
using a linear function. Alpha is the split frequency and Beta is the inclination of the line.
add everything under the line as the contribution of Primary and everything over the line as the contribution of Secondary
"""
if a.shape == ():
return alpha * a + (1 - alpha) * b
c_indices = torch.argsort(torch.flatten(c))
a_dist = torch.gather(torch.flatten(a), 0, c_indices)
b_dist = torch.gather(torch.flatten(b), 0, c_indices)
a_dft = torch.fft.rfft(a_dist.float())
b_dft = torch.fft.rfft(b_dist.float())
dft_filter = torch.arange(0, torch.numel(a_dft), device=a_dft.device).float()
dft_filter /= torch.numel(a_dft)
if beta > EPSILON:
dft_filter = (dft_filter - alpha) / beta + 1 / 2
dft_filter = torch.clamp(dft_filter, 0.0, 1.0)
else:
dft_filter = (dft_filter >= alpha).float()
x_dft = (1 - dft_filter) * a_dft + dft_filter * b_dft
x_dist = torch.fft.irfft(x_dft, a_dist.shape[0])
x_values = torch.gather(x_dist, 0, torch.argsort(c_indices))
return x_values.reshape_as(a)
def ties_add_difference(a: Tensor, b: Tensor, c: Tensor, alpha: float, beta: float, **kwargs) -> Tensor:
"""
An implementation of arXiv:2306.01708
"""
deltas = []
signs = []
for m in [a, b]:
deltas.append(filter_top_k(m - c, beta))
signs.append(torch.sign(deltas[-1]))
signs = torch.stack(signs, dim=0)
final_sign = torch.sign(torch.sum(signs, dim=0))
delta_filters = (signs == final_sign).float()
res = torch.zeros_like(c, device=c.device)
for delta_filter, delta in zip(delta_filters, deltas, strict=False):
res += delta_filter * delta
param_count = torch.sum(delta_filters, dim=0)
return c + alpha * torch.nan_to_num(res / param_count)
def filter_top_k(a: Tensor, k: float):
k = max(int((1 - k) * torch.numel(a)), 1)
k_value, _ = torch.kthvalue(torch.abs(a.flatten()).float(), k)
top_k_filter = (torch.abs(a) >= k_value).float()
return a * top_k_filter