import os import jax import time import copy import jax.nn as nn import jax.lax as lax import jax.numpy as jnp from lmc_model import print_model from collections import defaultdict from typing import NamedTuple from flax.core import freeze, unfreeze from jax import random, tree_util, jit, grad, value_and_grad from scipy.optimize import linear_sum_assignment, minimize import numpy as np import matplotlib.pyplot as plt def compute_objective(A, X, X_prime, Y, Y_prime): A_inv = np.linalg.inv(A) term1 = X - X_prime @ A.T term2 = Y - Y_prime @ A_inv return np.sum(term1**2) + np.sum(term2**2) def compute_gradient(A, X, X_prime, Y, Y_prime): A_inv = np.linalg.inv(A) term1 = -2 * X.T @ X_prime + 2 * A @ X_prime.T @ X_prime term2 = 2 * A_inv.T @ Y_prime.T @ (Y - Y_prime @ A_inv) @ A_inv.T return term1 + term2 def line_search(A, grad, X, X_prime, Y, Y_prime, max_step=1, tau=0.5, c1=1e-4): eta = max_step f_current = compute_objective(A, X, X_prime, Y, Y_prime) grad_norm2 = np.sum(grad**2) n = A.shape[0] while eta > 1e-10: A_new = A - eta * grad if np.linalg.matrix_rank(A_new) < n: eta *= tau continue f_new = compute_objective(A_new, X, X_prime, Y, Y_prime) if f_new <= f_current - c1 * eta * grad_norm2: return eta eta *= tau return 0 @jax.jit def compute_objective_jax(A, X, X_prime, Y, Y_prime, cond_threshold=1e6): cond = jnp.linalg.cond(A) def safe_obj(): A_inv = jnp.linalg.inv(A) term1 = X - X_prime @ A.T term2 = Y - Y_prime @ A_inv return jnp.sum(term1**2) + jnp.sum(term2**2) return lax.cond(cond > cond_threshold, lambda: jnp.inf, safe_obj) compute_value_and_grad_jax = jit(value_and_grad(compute_objective_jax)) def solve_orthogonal(X, X_prime, Y, Y_prime): B = X.T @ X_prime + Y.T @ Y_prime U, _, Vt = np.linalg.svd(B) return U @ Vt def optimize_alignment(A_init, X, X_prime, Y, Y_prime, max_iter=5000): objective_values = [] grad_norms = [] condition_nums = [] def obj_fn(flat_A): A = flat_A.reshape(A_init.shape) obj, grad_val = compute_value_and_grad_jax(jnp.array(A), jnp.array(X), jnp.array(X_prime), jnp.array(Y), jnp.array(Y_prime)) return float(obj), np.array(grad_val).flatten() def callback(flat_A): A = flat_A.reshape(A_init.shape) obj, grad_val = compute_value_and_grad_jax(jnp.array(A), jnp.array(X), jnp.array(X_prime), jnp.array(Y), jnp.array(Y_prime)) grad_norm = jnp.linalg.norm(grad_val, 'fro') cond = jnp.linalg.cond(jnp.array(A)) objective_values.append(float(obj)) grad_norms.append(float(grad_norm)) condition_nums.append(float(cond)) res = minimize(obj_fn, A_init.flatten(), jac=True, method='L-BFGS-B', options={'maxiter': max_iter}, callback=callback) A_opt = res.x.reshape(A_init.shape) return A_opt, objective_values, grad_norms, condition_nums def extract_attention_params(attn): key =np.array(attn['attention']['key']['kernel']) key_bias = np.array(attn['attention']['key']['bias']) query = np.array(attn['attention']['query']['kernel']) query_bias = np.array(attn['attention']['query']['bias']) value = np.array(attn['attention']['value']['kernel']) value_bias = np.array(attn['attention']['value']['bias']) out = np.array(attn['output']['dense']['kernel']) out_bias = np.array(attn['output']['dense']['bias']) return query, key, value, query_bias, key_bias, value_bias, out, out_bias def reshape_attention_weights(query, key, value, query_bias, key_bias, value_bias, out_kernel, num_heads): D = query.shape[0] D_k = D_v = D // num_heads def stack_per_head(tensor, axis=0): return np.stack([ tensor[:,i * D_k:(i + 1) * D_k] if axis == 0 else tensor[i * D_k:(i + 1) * D_k,:] for i in range(num_heads) ]) def stack_bias_per_head(bias): return np.stack([bias[i * D_k:(i + 1) * D_k] for i in range(num_heads)]) W_Q = stack_per_head(query) W_K = stack_per_head(key) W_V = stack_per_head(value) W_O = stack_per_head(out_kernel, axis=1) b_Q = stack_bias_per_head(query_bias) b_K = stack_bias_per_head(key_bias) b_V = stack_bias_per_head(value_bias) return W_Q, b_Q, W_K, b_K, W_V, b_V, W_O def compute_extended_weights(W, b): return np.vstack([W, b.reshape(1, -1)]) def compute_cost_matrix(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, h, activations, alpha=0.5): C = np.zeros((h, h)) for i in range(h): tilde_W_Q_a_i = np.vstack([W_Q_a[i], b_Q_a[i].reshape(1, -1)]) tilde_W_K_a_i = np.vstack([W_K_a[i], b_K_a[i].reshape(1, -1)]) tilde_W_V_a_i = np.vstack([W_V_a[i], b_V_a[i].reshape(1, -1)]) QKT_a_i = tilde_W_Q_a_i @ tilde_W_K_a_i.T VO_a_i = tilde_W_V_a_i @ W_O_a[i] centered_QKT_a_i = QKT_a_i - np.mean(QKT_a_i, axis=1, keepdims=True) for j in range(h): tilde_W_Q_b_j = np.vstack([W_Q_b[j], b_Q_b[j].reshape(1, -1)]) tilde_W_K_b_j = np.vstack([W_K_b[j], b_K_b[j].reshape(1, -1)]) tilde_W_V_b_j = np.vstack([W_V_b[j], b_V_b[j].reshape(1, -1)]) QKT_b_j = tilde_W_Q_b_j @ tilde_W_K_b_j.T VO_b_j = tilde_W_V_b_j @ W_O_b[j] centered_QKT_b_j = QKT_b_j - np.mean(QKT_b_j, axis=1, keepdims=True) cost = 0.5 * np.sum((centered_QKT_a_i - centered_QKT_b_j) ** 2) cost += 0.5 * np.sum((VO_a_i - VO_b_j) ** 2) C[i, j] = cost return C # def compute_cost_matrix(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, # W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, # num_heads, activations, alpha=0.5): # # activations: (B, L, D) # B, L, D = activations.shape # print("Activations shape: ", activations.shape) # # Augment activations with a column of ones for bias calculation # ones_col = jnp.ones((B, L, 1)) # X_tilde = jnp.concatenate([activations, ones_col], axis=-1) # Shape (B, L, D+1) # d_head = W_Q_a[0].shape[1] # sqrt_d = jnp.sqrt(float(d_head)) # C = np.zeros((num_heads, num_heads)) # for i in range(num_heads): # # Pre-compute for model A, head i # tilde_W_Q_a_i = compute_extended_weights(W_Q_a[i], b_Q_a[i]) # tilde_W_K_a_i = compute_extended_weights(W_K_a[i], b_K_a[i]) # tilde_W_V_a_i = compute_extended_weights(W_V_a[i], b_V_a[i]) # Q_a_i = X_tilde @ tilde_W_Q_a_i # K_a_i = X_tilde @ tilde_W_K_a_i # S_a_i = jnp.einsum('bld,bmd->blm', Q_a_i, K_a_i) / sqrt_d # P_a_i = nn.softmax(S_a_i, axis=-1) # V_tilde_a_i = X_tilde @ tilde_W_V_a_i # V_a_i = V_tilde_a_i @ W_O_a[i] # for j in range(num_heads): # # Compute for model B, head j # tilde_W_Q_b_j = compute_extended_weights(W_Q_b[j], b_Q_b[j]) # tilde_W_K_b_j = compute_extended_weights(W_K_b[j], b_K_b[j]) # tilde_W_V_b_j = compute_extended_weights(W_V_b[j], b_V_b[j]) # Q_b_j = X_tilde @ tilde_W_Q_b_j # K_b_j = X_tilde @ tilde_W_K_b_j # S_b_j = jnp.einsum('bld,bmd->blm', Q_b_j, K_b_j) / sqrt_d # P_b_j = nn.softmax(S_b_j, axis=-1) # V_tilde_b_j = X_tilde @ tilde_W_V_b_j # V_b_j = V_tilde_b_j @ W_O_b[j] # cost_P = jnp.sum((P_a_i - P_b_j)**2) # cost_V = jnp.sum((V_a_i - V_b_j)**2) # total_cost = (alpha * cost_P + (1 - alpha) * cost_V) / B # C[i, j] = total_cost # return C # def compute_cost_matrix(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, # W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, # num_heads, activations, alpha=0.5, epsilon=1e-8): # """ # Computes the cost matrix for attention head permutation using Cosine Similarity. # The cost C[i, j] is defined as 1 - cosine_similarity, which is minimized when # the output tensors are most similar. # """ # # activations: (B, L, D) # B, L, D = activations.shape # print("Activations shape: ", activations.shape) # # Augment activations with a column of ones for bias calculation # ones_col = jnp.ones((B, L, 1)) # X_tilde = jnp.concatenate([activations, ones_col], axis=-1) # Shape (B, L, D+1) # d_head = W_Q_a[0].shape[1] # sqrt_d = jnp.sqrt(float(d_head)) # C = np.zeros((num_heads, num_heads)) # for i in range(num_heads): # # Pre-compute for model A, head i # tilde_W_Q_a_i = compute_extended_weights(W_Q_a[i], b_Q_a[i]) # tilde_W_K_a_i = compute_extended_weights(W_K_a[i], b_K_a[i]) # tilde_W_V_a_i = compute_extended_weights(W_V_a[i], b_V_a[i]) # Q_a_i = X_tilde @ tilde_W_Q_a_i # K_a_i = X_tilde @ tilde_W_K_a_i # S_a_i = jnp.einsum('bld,bmd->blm', Q_a_i, K_a_i) / sqrt_d # P_a_i = nn.softmax(S_a_i, axis=-1) # P_a_i_flat = P_a_i.flatten() # V_tilde_a_i = X_tilde @ tilde_W_V_a_i # V_a_i = V_tilde_a_i @ W_O_a[i] # V_a_i_flat = V_a_i.flatten() # for j in range(num_heads): # # Compute for model B, head j # tilde_W_Q_b_j = compute_extended_weights(W_Q_b[j], b_Q_b[j]) # tilde_W_K_b_j = compute_extended_weights(W_K_b[j], b_K_b[j]) # tilde_W_V_b_j = compute_extended_weights(W_V_b[j], b_V_b[j]) # Q_b_j = X_tilde @ tilde_W_Q_b_j # K_b_j = X_tilde @ tilde_W_K_b_j # S_b_j = jnp.einsum('bld,bmd->blm', Q_b_j, K_b_j) / sqrt_d # P_b_j = nn.softmax(S_b_j, axis=-1) # P_b_j_flat = P_b_j.flatten() # V_tilde_b_j = X_tilde @ tilde_W_V_b_j # V_b_j = V_tilde_b_j @ W_O_b[j] # V_b_j_flat = V_b_j.flatten() # # Cosine similarity for P # dot_P = jnp.dot(P_a_i_flat, P_b_j_flat) # norm_P_a = jnp.linalg.norm(P_a_i_flat) # norm_P_b = jnp.linalg.norm(P_b_j_flat) # cos_sim_P = dot_P / (norm_P_a * norm_P_b + epsilon) # cost_P = 1.0 - cos_sim_P # # Cosine similarity for V # dot_V = jnp.dot(V_a_i_flat, V_b_j_flat) # norm_V_a = jnp.linalg.norm(V_a_i_flat) # norm_V_b = jnp.linalg.norm(V_b_j_flat) # cos_sim_V = dot_V / (norm_V_a * norm_V_b + epsilon) # cost_V = 1.0 - cos_sim_V # total_cost = (alpha * cost_P + (1 - alpha) * cost_V) # C[i, j] = total_cost # return C # Version 3: Normalized Norm # def compute_cost_matrix(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, # W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, # num_heads, activations, alpha=0.5, epsilon=1e-8): # """ # Computes the cost matrix for attention head permutation using Normalized Frobenius Norm. # The cost C[i, j] is the relative error: ||A-B||_F / ||A||_F. # """ # # activations: (B, L, D) # B, L, D = activations.shape # print("Activations shape: ", activations.shape) # # Augment activations with a column of ones for bias calculation # ones_col = jnp.ones((B, L, 1)) # X_tilde = jnp.concatenate([activations, ones_col], axis=-1) # Shape (B, L, D+1) # d_head = W_Q_a[0].shape[1] # sqrt_d = jnp.sqrt(float(d_head)) # C = np.zeros((num_heads, num_heads)) # for i in range(num_heads): # # Pre-compute for model A, head i # tilde_W_Q_a_i = compute_extended_weights(W_Q_a[i], b_Q_a[i]) # tilde_W_K_a_i = compute_extended_weights(W_K_a[i], b_K_a[i]) # tilde_W_V_a_i = compute_extended_weights(W_V_a[i], b_V_a[i]) # Q_a_i = X_tilde @ tilde_W_Q_a_i # K_a_i = X_tilde @ tilde_W_K_a_i # S_a_i = jnp.einsum('bld,bmd->blm', Q_a_i, K_a_i) / sqrt_d # P_a_i = nn.softmax(S_a_i, axis=-1) # V_tilde_a_i = X_tilde @ tilde_W_V_a_i # V_a_i = V_tilde_a_i @ W_O_a[i] # for j in range(num_heads): # # Compute for model B, head j # tilde_W_Q_b_j = compute_extended_weights(W_Q_b[j], b_Q_b[j]) # tilde_W_K_b_j = compute_extended_weights(W_K_b[j], b_K_b[j]) # tilde_W_V_b_j = compute_extended_weights(W_V_b[j], b_V_b[j]) # Q_b_j = X_tilde @ tilde_W_Q_b_j # K_b_j = X_tilde @ tilde_W_K_b_j # S_b_j = jnp.einsum('bld,bmd->blm', Q_b_j, K_b_j) / sqrt_d # P_b_j = nn.softmax(S_b_j, axis=-1) # V_tilde_b_j = X_tilde @ tilde_W_V_b_j # V_b_j = V_tilde_b_j @ W_O_b[j] # # Normalized Frobenius norm for P. We flatten the 3D tensor to compute the L2 norm. # norm_diff_P = jnp.linalg.norm((P_a_i - P_b_j).flatten()) # norm_P_a = jnp.linalg.norm(P_a_i.flatten()) # cost_P = norm_diff_P / (norm_P_a + epsilon) # # Normalized Frobenius norm for V. We flatten the 3D tensor to compute the L2 norm. # norm_diff_V = jnp.linalg.norm((V_a_i - V_b_j).flatten()) # norm_V_a = jnp.linalg.norm(V_a_i.flatten()) # cost_V = norm_diff_V / (norm_V_a + epsilon) # total_cost = (alpha * cost_P + (1 - alpha) * cost_V) # C[i, j] = total_cost # return C def align_single_head(W_Q_a_i, b_Q_a_i, W_K_a_i, b_K_a_i, W_V_a_i, b_V_a_i, W_O_a_i, W_Q_b_i, b_Q_b_i, W_K_b_i, b_K_b_i, W_V_b_i, b_V_b_i, W_O_b_i, optimize): tilde_W_Q_a_i = compute_extended_weights(W_Q_a_i, b_Q_a_i) tilde_W_K_a_i = compute_extended_weights(W_K_a_i, b_K_a_i) tilde_W_V_a_i = compute_extended_weights(W_V_a_i, b_V_a_i) Y_O_a_i = W_O_a_i.T tilde_W_Q_b_i = compute_extended_weights(W_Q_b_i, b_Q_b_i) tilde_W_K_b_i = compute_extended_weights(W_K_b_i, b_K_b_i) tilde_W_V_b_i = compute_extended_weights(W_V_b_i, b_V_b_i) Y_O_b_i = W_O_b_i.T A_init = solve_orthogonal(tilde_W_Q_a_i, tilde_W_Q_b_i, tilde_W_K_a_i, tilde_W_K_b_i) B_init = solve_orthogonal(Y_O_a_i, Y_O_b_i, tilde_W_V_a_i, tilde_W_V_b_i) if optimize: A, objective_values_A, grad_norms_A, condition_nums_A = optimize_alignment( A_init, tilde_W_Q_a_i, tilde_W_Q_b_i, tilde_W_K_a_i, tilde_W_K_b_i ) B, objective_values_B, grad_norms_B, condition_nums_B = optimize_alignment( B_init, Y_O_a_i, Y_O_b_i, tilde_W_V_a_i, tilde_W_V_b_i ) else: A = A_init B = B_init A_inv = np.linalg.inv(A) B_inv = np.linalg.inv(B) W_Q_aligned = W_Q_b_i @ A.T b_Q_aligned = b_Q_b_i @ A.T W_K_aligned = W_K_b_i @ A_inv b_K_aligned = b_K_b_i @ A_inv W_V_aligned = W_V_b_i @ B_inv b_V_aligned = b_V_b_i @ B_inv W_O_aligned = B @ W_O_b_i aligned_params = { 'query': {'kernel': W_Q_aligned, 'bias': b_Q_aligned}, 'key': {'kernel': W_K_aligned, 'bias': b_K_aligned}, 'value': {'kernel': W_V_aligned, 'bias': b_V_aligned}, 'out': {'kernel': W_O_aligned} } if optimize: return { 'aligned_params': aligned_params, 'metrics_A': { 'objective_values': objective_values_A, 'grad_norms': grad_norms_A, 'condition_nums': condition_nums_A }, 'metrics_B': { 'objective_values': objective_values_B, 'grad_norms': grad_norms_B, 'condition_nums': condition_nums_B } } return {'aligned_params': aligned_params} def apply_alignment(W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, h): aligned_params = {} for i in range(h): print(f"Aligning Heads {i}") tilde_W_Q_a_i = np.vstack([W_Q_a[i], b_Q_a[i].reshape(1, -1)]) tilde_W_K_a_i = np.vstack([W_K_a[i], b_K_a[i].reshape(1, -1)]) tilde_W_Q_b_i = np.vstack([W_Q_b[i], b_Q_b[i].reshape(1, -1)]) tilde_W_K_b_i = np.vstack([W_K_b[i], b_K_b[i].reshape(1, -1)]) A_i_init = solve_orthogonal(tilde_W_Q_a_i, tilde_W_Q_b_i, tilde_W_K_a_i, tilde_W_K_b_i) # A_i = optimize(A_i_init, tilde_W_Q_a_i, tilde_W_Q_b_i, tilde_W_K_a_i, tilde_W_K_b_i) A_i = A_i_init tilde_W_V_a_i = np.vstack([W_V_a[i], b_V_a[i].reshape(1, -1)]) tilde_W_V_b_i = np.vstack([W_V_b[i], b_V_b[i].reshape(1, -1)]) Y_O = W_O_a[i].T Y_O_prime = W_O_b[i].T B_i_init = solve_orthogonal(tilde_W_V_a_i, tilde_W_V_b_i, Y_O, Y_O_prime) # B_i = optimize(B_i_init, tilde_W_V_a_i, tilde_W_V_b_i, Y_O, Y_O_prime) B_i = B_i_init A_i_inv = np.linalg.inv(A_i) B_i_inv = np.linalg.inv(B_i) W_Q_aligned = W_Q_b[i] @ A_i.T b_Q_aligned = b_Q_b[i] @ A_i.T W_K_aligned = W_K_b[i] @ A_i_inv b_K_aligned = b_K_b[i] @ A_i_inv W_V_aligned = W_V_b[i] @ B_i_inv b_V_aligned = b_V_b[i] @ B_i_inv W_O_aligned = B_i @ W_O_b[i] aligned_params[f'head_{i}'] = { 'query': {'kernel': W_Q_aligned, 'bias': b_Q_aligned}, 'key': {'kernel': W_K_aligned, 'bias': b_K_aligned}, 'value': {'kernel': W_V_aligned, 'bias': b_V_aligned}, 'out': {'kernel': W_O_aligned} } return aligned_params def merge_aligned_params(aligned_params, h, D, out_bias_b): query_kernel = np.stack([aligned_params[f'head_{i}']['query']['kernel'] for i in range(h)], axis=1) query_bias = np.stack([aligned_params[f'head_{i}']['query']['bias'] for i in range(h)], axis=0) key_kernel = np.stack([aligned_params[f'head_{i}']['key']['kernel'] for i in range(h)], axis=1) key_bias = np.stack([aligned_params[f'head_{i}']['key']['bias'] for i in range(h)], axis=0) value_kernel = np.stack([aligned_params[f'head_{i}']['value']['kernel'] for i in range(h)], axis=1) value_bias = np.stack([aligned_params[f'head_{i}']['value']['bias'] for i in range(h)], axis=0) out_kernel = np.stack([aligned_params[f'head_{i}']['out']['kernel'] for i in range(h)], axis=0) # print(query_kernel.shape,"|",query_bias.shape) # print(key_kernel.shape,"|",key_bias.shape) # print(value_kernel.shape,"|",value_bias.shape) # print(out_kernel.shape,) query_kernel = query_kernel.reshape(-1, D) key_kernel = key_kernel.reshape(-1, D) value_kernel = value_kernel.reshape(-1, D) out_kernel = out_kernel.reshape(D, -1) query_bias = query_bias.reshape(-1) key_bias = key_bias.reshape(-1) value_bias = value_bias.reshape(-1) return { 'attention': { 'query': {'kernel': jnp.array(query_kernel),'bias': jnp.array(query_bias),}, 'value': {'kernel': jnp.array(value_kernel),'bias': jnp.array(value_bias),}, 'key': {'kernel': jnp.array(key_kernel),'bias': jnp.array(key_bias),}, }, 'output': { 'dense': {'kernel': jnp.array(out_kernel),'bias': jnp.array(out_bias_b),}, }, } def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activation, permute_heads=True, optimize=False, alpha=0.5): attn_a = params_a['vit']['encoder']["layer"][str(layer_idx)]['attention'] attn_b = params_b['vit']['encoder']["layer"][str(layer_idx)]['attention'] query_a, key_a, value_a, query_bias_a, key_bias_a, value_bias_a, out_a, out_bias_a = extract_attention_params(attn_a) query_b, key_b, value_b, query_bias_b, key_bias_b, value_bias_b, out_b, out_bias_b = extract_attention_params(attn_b) W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a = reshape_attention_weights(query_a, key_a, value_a, query_bias_a, key_bias_a, value_bias_a, out_a, num_heads) W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b = reshape_attention_weights(query_b, key_b, value_b, query_bias_b, key_bias_b, value_bias_b, out_b, num_heads) if permute_heads: C = compute_cost_matrix(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, num_heads, activation, alpha) row_ind, col_ind = linear_sum_assignment(C) print("Best Permutation Heads:", col_ind) W_Q_b = [W_Q_b[j] for j in col_ind] b_Q_b = [b_Q_b[j] for j in col_ind] W_K_b = [W_K_b[j] for j in col_ind] b_K_b = [b_K_b[j] for j in col_ind] W_V_b = [W_V_b[j] for j in col_ind] b_V_b = [b_V_b[j] for j in col_ind] W_O_b = [W_O_b[j] for j in col_ind] if optimize: metrics_A_all = {key: [] for key in ['objective_values', 'grad_norms', 'condition_nums']} metrics_B_all = {key: [] for key in ['objective_values', 'grad_norms', 'condition_nums']} aligned_params, return_dict = {}, {} for i in range(num_heads): result = align_single_head( W_Q_a[i], b_Q_a[i], W_K_a[i], b_K_a[i], W_V_a[i], b_V_a[i], W_O_a[i], W_Q_b[i], b_Q_b[i], W_K_b[i], b_K_b[i], W_V_b[i], b_V_b[i], W_O_b[i], optimize ) aligned_params[f'head_{i}'] = result['aligned_params'] if optimize: for key in metrics_A_all: metrics_A_all[key].append(result['metrics_A'][key]) metrics_B_all[key].append(result['metrics_B'][key]) return_dict['aligned_params'] = merge_aligned_params(aligned_params, num_heads, query_a.shape[1], out_bias_b) if optimize: return_dict['metrics_A_all'] = metrics_A_all return_dict['metrics_B_all'] = metrics_B_all return return_dict def weight_matching_attn(rng, params_a, params_b, activation, config): params_dict = {} configurations = [ # ("permu_head_init_ortho_no_opt", 'ortho', True, False), ("permu_head_init_ortho_opt", 'ortho', True, True), # ("naive_head_init_ortho_no_opt", 'ortho', False, False), # ("naive_head_init_ortho_opt", 'ortho', False, True), ] for name, init_method, permute_heads, optimize in configurations: aligned_params = copy.deepcopy(params_b) if optimize: layer_to_metrics_A = {} layer_to_metrics_B = {} for layer_idx in config.lmc_layer_indices: if activation is not None: activations_for_layer = activation[layer_idx] else: activations_for_layer = None result = align_attention_params( rng, params_a, aligned_params, layer_idx, config.lmc_config.num_attention_heads, activations_for_layer, permute_heads=permute_heads, optimize=optimize ) aligned_params['vit']['encoder']["layer"][str(layer_idx)]['attention'] = result['aligned_params'] if optimize: layer_to_metrics_A[layer_idx] = result['metrics_A_all'] layer_to_metrics_B[layer_idx] = result['metrics_B_all'] total_sum = tree_util.tree_reduce(lambda acc, x: acc + jnp.sum(x), aligned_params, initializer=0) print(f"{name}: {total_sum}, sanity check") params_dict[name] = aligned_params return params_dict # cost_head = copy.deepcopy(params_b) # naive_head = copy.deepcopy(params_b) # for layer_idx in config.lmc_layer_indices: # aligned_attention_params = cost_align_attention_params(rng, params_a, params_b, str(layer_idx), config.lmc_config.num_attention_heads) # cost_head['vit']['encoder']["layer"][str(layer_idx)]['attention'] = aligned_attention_params # for layer_idx in config.lmc_layer_indices: # aligned_attention_params = naive_align_attention_params(rng, params_a, params_b, str(layer_idx), config.lmc_config.num_attention_heads) # naive_head['vit']['encoder']["layer"][str(layer_idx)]['attention'] = aligned_attention_params # return {"cost_head": cost_head, "naive_head": naive_head}