<jupyter_start><jupyter_text>Summary Author: <NAME>Learning a projection matrix for CIFAR and MNIST datasets *but also training $\beta$. Results here are inconclusive in that we need to find optimal $d^*$ for these larger and correlated $m$ patterns and how these are fit by $\beta$. In some cases they also seem to be very large, other times smaller but this is also function of their initailization values and the size of the perturbations used during training. Ultimately these results and not as relevant as relating to actual pretrained Transformer models.<jupyter_code>%load_ext autoreload
%autoreload 2
from Implementations_Associative_Memory import *
from Data_Processing_Associative_Memory import *
from utils_LearnProjections import *
from SDM_Circ_Inter_Funcs import *
import pickle
import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import scipy
import torch.nn.functional as F
import copy
from types import SimpleNamespace
torch.cuda.set_device(2) 
data_name = "MNIST" #"MNIST" # "CIFAR"
load_in_W_k = False
learn_beta = True
n_train_repeats = 3 # number of weight matrices trained from random initialization for each parameter setting.
n_perturb_repeats = 1 # number of times perturbations are applied to the same learned weight matrix.
if learn_beta:
    neuron_vals = [None]
else: 
    neuron_vals = list(np.round(np.logspace(6,20,5))) # testing the algorithms with different numbers of neurons
    neuron_vals.append(None)
hamm_radii = [5, 9, 11, 15] 
max_perturb = 300
n_perturbs = 7 #number of perturbations to apply. 
cont_circle_intersect = False # fitting beta to binary SDM during training using the Attention softmax

m=None #None means to use the full database. 
batch_size = 128
n=64
r = 2**n
tbatch_size = 2000
num_workers = 4
print_each_result = False 
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Device is:", device)
algorithms = ["Cont_SDM_CFit_Attention", "Cont_SDM", "Cont_Binary_SDM", "Cont_SDM_BFit_Attention"] 
train_test_iter = [True, False] # evaluates trained projection matrices on train and test datasets
rad = 1 # circle radius size. Unit norm here. 

# None values are placeholders. 
sdm_params = dict(
    n=n, 
    n_base = None, # used as the base dimension, 
    r=r, 
    m=m,
    batch_size = batch_size,
    hamm_radius =None,
    fix_perturb=None,
    proj_epochs = 20, # number of epochs to train weight matrix. 10 gives good convergence.
    lr=0.01,
    max_conv_steps = 10, # steps to update the query
    epsilon = 0.05, # removes any rounding error in computing the similarity.
    ensure_no_closer_images = False,
    min_train_perturb = 10,
    max_train_perturb = 100,
    perturb=None,
    device=device,
    cont_cache_resolution = 10000, # number of points to use when computing the continuous circle intersection
    plot_batches = None, # plots losses during training more often
    plot_epochs = False,
    binary_data = False, # only using continuous vectors here.
    renorm_cont_queries = False, # only want to L2 norm after projecting, not on the raw images that are non L2 normed.
    train_multiple_steps = False, # allow for multiple steps during training. 
    # that have changed by less than epsilon at any of their pixel locations
    learn_beta = learn_beta
)
sdm_params = SimpleNamespace(**sdm_params) # enabling use of dot notation on the dictionary

# Generating datasets
if data_name=="MNIST":
    image_len = 28
    transform = transforms.Compose(
            [transforms.ToTensor()])
    trainset = torchvision.datasets.MNIST(root='./data', train=True,
                                            download=True, transform=transform)
    
    
    testset = torchvision.datasets.MNIST(root='./data', train=False,
                                           download=True, transform=transform)

elif data_name=="CIFAR":
    image_len = 32
    transform = transforms.Compose(
        [transforms.Grayscale(num_output_channels=1),
        transforms.ToTensor()
        ])

    trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                            download=True, transform=transform)
    
    testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                           download=True, transform=transform)
    classes = ('plane', 'car', 'bird', 'cat',
               'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
    
print("Train and Test sets",trainset, testset)
loader_size_train = len(trainset) if m is None else m
loader_size_test = len(testset) if m is None else m
trainloader = torch.utils.data.DataLoader(trainset, batch_size=loader_size_train,
                                      shuffle=True, num_workers=num_workers)
testloader = torch.utils.data.DataLoader(testset, batch_size=loader_size_test,
                                      shuffle=True, num_workers=num_workers)
# full image dimensionality, different from n=64 in the latent space projected into.
n_base = image_len**2
sdm_params.n_base = n_base

# set the loaders to give only batches. 
# somewhat hacky as using the loader from before to serve the whole dataset. 
dataiter = iter(trainloader)
keys, labels = dataiter.next()
nchannels = keys.shape[1]
# show example images:
imshow(torchvision.utils.make_grid(keys[:batch_size].view(batch_size, nchannels, image_len, image_len)))
cpu_keys = keys.flatten(start_dim=1)
trainloader = torch.utils.data.DataLoader(cpu_keys, batch_size=batch_size,
                                          shuffle=True, num_workers=num_workers)
keys = torch.clone(cpu_keys).to(device).T

# setting up test loader
dataiter = iter(testloader)
tkeys, labels = dataiter.next()
print(tkeys.shape)
cpu_tkeys = tkeys.flatten(start_dim=1)
testloader = torch.utils.data.DataLoader(cpu_tkeys, batch_size=tbatch_size,
                                         shuffle=True, num_workers=num_workers)
tkeys = torch.clone(cpu_keys).to(device).T

perturb_hamms = list(np.unique(np.linspace(0,max_perturb,n_perturbs).astype(int)))
perturb_hamms
"""
Iterate through hamming distances. For n_train_repeats train the projection matrix. 
Then iterate through the different algorithms to compare them. 
"""

neuron_res_nrepeats = {k:[] for k in neuron_vals}
hamm_res_nrepeats = {k:copy.deepcopy(neuron_res_nrepeats) for k in hamm_radii}
all_res = {k:copy.deepcopy(hamm_res_nrepeats) for k in algorithms}

for hamm_radius in hamm_radii:
    print('hamm_radius is', hamm_radius)
    sdm_params.hamm_radius = hamm_radius
    # varying r doesnt matter for this beta fit. For now fitting to the binary SDM beta value!
    
    

    for train_rep_ind in range(n_train_repeats):
        # reset W_k and beta!
        # fitting SDM beta here: 
        beta = fit_beta(sdm_params.n, sdm_params.r, sdm_params.hamm_radius, False, plot=False)
        W_k = torch.randn((n,n_base), requires_grad=True, device=device)
        if learn_beta:
            beta = torch.nn.Parameter(3*beta*torch.ones(1).to(device), requires_grad=True)
            beta_name = 'trained_weights/trained_beta'+data_name+"|m="+str(m)+'|'+str(hamm_radius)+'|'+str(train_rep_ind)+'.torch'
        
        sdm_params.fix_perturb=False
        W_k_name = 'trained_weights/trained_W_k'+data_name+'learnBeta'+str(learn_beta)+"|m="+str(m)+'|'+str(hamm_radius)+'|'+str(train_rep_ind)+'.torch'
        if load_in_W_k:
            W_k = torch.load(W_k_name).to(device)
            if learn_beta:
                beta = torch.load(beta_name).to(device)
        else: 
            trained_params = train_projection(W_k, sdm_params, trainloader, keys, beta)
            if learn_beta: 
                W_k, beta = trained_params
            else: 
                W_k = trained_params
            torch.save(W_k, W_k_name)
            torch.save(beta, beta_name)

        # testing everything.
        sdm_params.fix_perturb=True
        for pert_rep_ind in range(n_perturb_repeats):
        
            for algo in algorithms:
                if "Attention" in algo:
                    cos_res = test_projection(algo, perturb_hamms, W_k, 
                                                        sdm_params, trainloader, 
                                                        testloader, keys, 
                                                        tkeys,
                                                        None, beta, plot_results=False )
                    all_res[algo][hamm_radius][neuron_vals[-1]].append(cos_res)

                else: 
                    # set the neuron value
                    for r in neuron_vals:
                        sdm_params.r=r
                        
                        b_cached_inters, log_c_cached_inters = get_binary_and_continuous_caches(n, hamm_radius, r, 
                                                                                    sdm_params.cont_cache_resolution)
                        b_cached_inters = torch.Tensor(b_cached_inters).to(device)
                        log_c_cached_inters = torch.Tensor(log_c_cached_inters.astype(float)).to(device)
                        
                        use_binary_intersect = True if algo =="Cont_Binary_SDM" else False
                        input_cache = b_cached_inters if use_binary_intersect else log_c_cached_inters
                        
                        should_continue=check_caches_work_for_num_neurons(r, hamm_radius, use_binary_intersect, b_cached_inters, log_c_cached_inters)
                        if should_continue == 'continue':
                            continue
                            
                        cos_res = test_projection(algo, perturb_hamms, W_k, 
                                                    sdm_params, trainloader, 
                                                    testloader, keys,
                                                    tkeys, input_cache, beta, plot_results=False )
                        all_res[algo][hamm_radius][r].append(cos_res)


            with open(data_name+"|m="+str(m)+'learnBeta'+str(learn_beta)+"train_projection_runs.pkl", 'wb') as fname: 
                pickle.dump(all_res, fname)

for hamm_radius in hamm_radii:
    sdm_beta = fit_beta(sdm_params.n, sdm_params.r, hamm_radius, False, plot=False)
    print('sdm beta', sdm_beta, 'input', 3*sdm_beta)
    for train_rep_ind in range(n_train_repeats):
        beta_name = 'trained_weights/trained_beta'+data_name+"|m="+str(m)+'|'+str(hamm_radius)+'|'+str(train_rep_ind)+'.torch'
        beta = torch.load(beta_name).to(device)
        print(hamm_radius, beta)
all_res.keys()
load_all_res = True

if load_all_res:
    with open(data_name+"|m="+str(m)+"train_projection_runs.pkl", 'rb') as fname: 
        all_res = pickle.load(fname)<jupyter_output><empty_output><jupyter_text>Plots of all results:<jupyter_code>directory='figures/LearnProjections/'

nneurons = neuron_vals[-1]

for train in [True, False]:
    for key, res_per_key in all_res.items():
        # each is a dictionary i need to homogenize. 
        for hamm_radius in np.sort(hamm_radii):

            agg = []

            for res in res_per_key[hamm_radius][nneurons]:
                if train: 
                    agg.append(res['train'])
                    labl = 'Train'
                else: 
                    agg.append(res['test'])
                    labl='Test'
            agg = np.asarray(agg)
            print(agg.shape)
            iters_mu = agg.mean(axis=(0,2)) # mean for each iteration value
            iters_std = agg.std(axis=(0,2))

            #stds = np.stack([np.asarray(stds), np.zeros((len(stds)))])
            plt.errorbar(perturb_hamms, iters_mu, yerr=iters_std,capsize=2, label='$d$='+str(hamm_radius))
        
        plt.scatter(perturb_hamms, hamm_to_cosine(np.asarray(perturb_hamms),n_base), label='Baselines')
        
        if nneurons is None or nneurons == "None": 
            neuron_label = 'Infinite'
        else:
            print(float(nneurons))
            neuron_label = '$r$='+"{:.0e}".format(nneurons)
            
        if m is None: 
            printm= len(trainset)
            
        plt.title("Cosine Similarity across Perturbations"+"\n"+"Type: "+key+" | Dataset: "+data_name+" | $r$="+neuron_label+" | $m=$"+str(printm)+" | "+labl)
        plt.legend()
        plt.xlabel("Perturbation")
        plt.ylabel("Cosine Similarity")
        save_str = "Dataset="+data_name+'learnBeta'+str(learn_beta)+"|n="+str(n)+"m="+str(m)+"r="+str(r)+"Type="+key+"|NumNeuron="+neuron_label+"|"+labl
        #plt.gcf().savefig(directory+save_str+'.png', dpi=250)
        plt.show()<jupyter_output>(3, 7, 60000)
(3, 7, 60000)
(3, 7, 60000)
(3, 7, 60000)<jupyter_text>Plots varying the number of neurons within each Hamming radius setting<jupyter_code>cmap = plt.get_cmap("tab10") # used to keep colors of each line constant across plots

for train in [True]:
    for key, res_per_key in all_res.items():
        # each is a dictionary i need to homogenize. 
        for hamm_radius in np.sort(hamm_radii):
            for nind, nneurons in enumerate(neuron_vals):
                
                if len(res_per_key[hamm_radius][nneurons])==0:
                        print("skipping", hamm_radius, nneurons)
                        continue
                agg = []
                for res in res_per_key[hamm_radius][nneurons]:
                    if train: 
                        agg.append(res['train'])
                        labl = 'Train'
                    else: 
                        agg.append(res['test'])
                        labl='Test'
                        
                agg = np.asarray(agg)

                iters_mu = agg.mean(axis=(0,2)) # mean for each iteration value
                iters_std = agg.std(axis=(0,2))

                #stds = np.stack([np.asarray(stds), np.zeros((len(stds)))])
                print(hamm_radius, nneurons, iters_mu)
                if nneurons is None or nneurons == "None": 
                    neuron_label = 'Infinite'
                else:
                    print(float(nneurons))
                    neuron_label = '$r$='+"{:.0e}".format(nneurons)
                plt.errorbar(perturb_hamms, iters_mu, color=cmap(nind), yerr=iters_std,capsize=2, label=neuron_label) 
            plt.scatter(perturb_hamms, hamm_to_cosine(np.asarray(perturb_hamms),n_base), label='Baselines')
                
            plt.title("Cosine Similarity by Neuron Number"+" | Type: "+key+"\nDataset: "+data_name+" | $m$="+str(printm)+" | Hamm. Rad. $d$="+str(hamm_radius)+" | "+labl)
            plt.legend()
            plt.xlabel("Perturbation")
            plt.ylabel("Cosine Similarity")
            save_str = "NNeurons_Dataset="+data_name+'learnBeta'+str(learn_beta)+"|n="+str(n)+"m="+str(m)+"Type="+key+"|NumNeuron="+neuron_label+"|d="+str(hamm_radius)+"|"+labl
            plt.gcf().savefig(directory+save_str+'.png', dpi=250)
            plt.show()<jupyter_output>5 None [0.99839705 0.99810773 0.99695593 0.98543257 0.9226592  0.7475696
 0.54196763]<jupyter_text>Display convergence and perturbations:<jupyter_code>sdm_params
# baselines for the percentage reconstructed:

sdm_params.fix_perturb = True
sdm_params.stop_at_epsilon_convergence_per_datapoint = True 
sdm_params.max_conv_steps = 10
algorithms = ['Cont_SDM',
 'Cont_Binary_SDM',
 'Cont_SDM_BFit_Attention',
 'Cont_SDM_CFit_Attention']

load_in_W_k=True
hamm_radius=5
train_rep_ind=1
r=None#neuron_vals[-1]
train_o_test = 'train' # Test
# todo enable the test loader to work here
curr_loader = trainloader if train_o_test=='train' else testloader
keys_to_use = keys if train_o_test=='train' else tkeys
sdm_params.batch_size = batch_size if train_o_test=='train' else tbatch_size
nimages = 1
sdm_params.binary_data = False
binary_data = False

sdm_params.hamm_radius = hamm_radius
sdm_params.r = r

if load_in_W_k:
    W_k_name = 'trained_weights/trained_W_k'+data_name+"|m="+str(m)+'|'+str(hamm_radius)+'|'+str(train_rep_ind)+'.torch'
    print("load name is:", W_k_name)
    W_k = torch.load(W_k_name).to(device)

for batch_ind, queries in enumerate(curr_loader): 
    queries = queries.to(device).T
    targets = torch.clone(queries)
    norm_targets = targets/torch.norm(targets, dim=0,keepdim=True)
    break 
     
all_p_queries = None
with torch.no_grad():
    pKs = W_k@keys_to_use
    pKs = pKs/torch.norm(pKs, dim=0,keepdim=True)
    # set the neuron value
    sdm_params.r=r
    #print('=======================',r)
    # compute the caches
    b_cached_inters, log_c_cached_inters = get_binary_and_continuous_caches(n, hamm_radius, r, 
                                                                                    sdm_params.cont_cache_resolution)
    b_cached_inters = torch.Tensor(b_cached_inters).to(device)
    log_c_cached_inters = torch.Tensor(log_c_cached_inters.astype(float)).to(device)
                        
    print(hamm_to_cosine(np.asarray(perturb_hamms),n_base))
    
    for perturb in perturb_hamms:

        sdm_params.perturb = perturb
        perturbed_queries, isclosests = cosine_perturb_queries(queries, keys_to_use, sdm_params)
        
        for algo_ind, algo in enumerate(algorithms):
            
            use_binary_intersect = True if binary_data else False
            if algo =="Cont_Binary_SDM":
                use_binary_intersect = True # overrriding this. 
            input_cache = b_cached_inters if use_binary_intersect else log_c_cached_inters

            if "Attention" in algo:
                if "CFit" in algo:
                    cont_circle_intersect = True
                elif "BFit" in algo: 
                    cont_circle_intersect = False
                else: 
                    raise Exception("Attention but neither B or C fit in the Attention function name.")
                beta = fit_beta(sdm_params.n, 2**64, sdm_params.hamm_radius, cont_circle_intersect, plot=False)
            else: 
                beta = None
        
            to_converge = torch.clone(perturbed_queries)
            
            '''plt.imshow(perturbed_queries[:,:nimages].cpu().numpy().reshape(28,28))
            plt.title("perturbed")
            plt.show()
            
            plt.imshow(to_converge[:,:nimages].cpu().numpy().reshape(28,28))
            plt.title("Input into convergence")
            plt.show()'''
            
            #print(beta, input_cache, algo)
            
            converged_queries = torch_run_algo(algo, sdm_params, to_converge, pKs, keys_to_use,
                                                                 input_cache, beta, W_k=W_k)
            
            
            
            '''plt.imshow(converged_queries[:,:nimages].cpu().numpy().reshape(28,28))
            plt.title('convergece')
            plt.show()'''
            
            print(algo, perturb, 'with target', torch_cosine_sims(converged_queries[:,:nimages] , targets[:,:nimages] ), 
                  'with perturb', torch_cosine_sims(converged_queries[:,:nimages] , perturbed_queries[:,:nimages] ) )
            
            
            if all_p_queries is None:
                all_p_queries = perturbed_queries[:,:nimages].unsqueeze(0).cpu()
            elif algo_ind==0:
                all_p_queries = torch.cat( [all_p_queries, perturbed_queries[:,:nimages].unsqueeze(0).cpu()])

            all_p_queries = torch.cat( [all_p_queries, converged_queries[:,:nimages].unsqueeze(0).cpu()])

print( all_p_queries.shape )
all_p_queries = torch.transpose(all_p_queries,  1,2).flatten(end_dim=-2)
fig = plt.gcf()
fig.set_size_inches(18.5, 14)
randn = np.round(np.random.rand(1)*10000).astype(int)[0]
    
imshow(torchvision.utils.make_grid(all_p_queries.view(all_p_queries.shape[0], nchannels,image_len,image_len),
                                   nrow=nimages*5), show_now=False)
plt.title("Perturbations | Cont_SDM | Cont_Binary_SDM | Cont_SDM_BFit_Attention | Cont_SDM_CFit_Attention\n$d$="+str(hamm_radius)+" | $r=$Infinite | "+train_o_test)
#plt.ylabel("Increasing Perturbations")
plt.tight_layout()
fig.savefig('figures/'+data_name+train_o_test+str(randn)+'Convergences.png', dpi=100)
plt.show()<jupyter_output>Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).<jupyter_text>Misc. evaluation of the images learnt, not used in the paper.<jupyter_code>epsilon = 0.05
test_epsilon = 0.2 # debatable what this value should be. 
n_to_plot = 5
use_softmax=False
use_continuous = True
fix_perturb=False
check_for_closer_images = False
with torch.no_grad():
    for batch_ind, queries in enumerate(trainloader): 
        # get a minibatch
        bs = len(queries)
        print(queries.shape)
        queries = queries.to(device).T # making the queries into columns. 
        targets = torch.clone(queries)
        #project keys
        pKs = W_k@keys
        pKs = pKs/torch.norm(pKs, dim=0,keepdim=True)

        queries, isclosests, nsteps = convergence_loop(queries, norm_keys, keys)
        #print('nsteps', nsteps)
        # reconstruction loss after convergence iterations. 
        loss = ((queries - targets)**2).sum()

        perc_recon =  torch.logical_and(targets-test_epsilon < queries , 
                                 targets+test_epsilon > queries).sum(0).cpu().detach().numpy() / img_flat_len  

        agree_spots =torch.logical_and(targets-test_epsilon < queries , 
                                     targets+test_epsilon > queries).cpu().detach().numpy()

        for ind in range(bs):

            query, target, agree_spot = queries[:, ind].cpu().detach(), targets[:,ind].cpu().detach(), agree_spots[:,ind]

            plt.imshow(agree_spot.reshape(image_len,image_len))
            print( "perc recon", perc_recon[ind] )
            print( 'num agree spots', agree_spot.sum(), img_flat_len)
            plt.show()

            ct = torch.cat( [target.unsqueeze(0), query.unsqueeze(0)] )
            imshow(torchvision.utils.make_grid(ct.view(2, nchannels,image_len,image_len)))

            #imshow(target.detach().view(nchannels,image_len,image_len))
            #imshow(query.detach().view(nchannels,image_len,image_len))

            print("===========")
            if ind> n_to_plot:
                break
        break
plt.hist((norm_keys@query/torch.norm(query)).detach().numpy(), bins=100)
plt.show()

pQ = W_k@query
# L2 norm all
pQ = pQ/torch.norm(pQ)
# columnwise cosine similarities
dotp = pKs.T@pQ

plt.hist(dotp.detach().numpy(), bins=100)
plt.show()
tp=torch.cat((target.unsqueeze(0),query.unsqueeze(0)))
imshow(torchvision.utils.make_grid(tp.view(2, nchannels, image_len, image_len).detach() ))
((query - target)**2).sum()
target.detach().view(nchannels,image_len, image_len);
query.detach().view(nchannels,image_len, image_len);
imshow(((query - target)**2).detach().view(nchannels,image_len, image_len))
torch.max(norm_keys@queries, dim=0)<jupyter_output><empty_output><jupyter_text>Validating the query perturbation is correct<jupyter_code>dataiter = iter(trainloader)
queries = dataiter.next()
queries = queries.to(sdm_params.device).T
queries.shape
sdm_params.fix_perturb = True
sdm_params.perturb=50

pqueries, isclosests = perturb_queries(queries, sdm_params, norm_keys)
pqueries.shape
hamm_to_cosine(50,img_flat_len)
(pqueries[:,0]/torch.norm(pqueries[:,0])).T @ (queries[:,0]/torch.norm(queries[:,0]))
pqueries[:,0].T@pqueries[:,0]
ct = torch.cat( [queries[:,0].cpu().unsqueeze(0), pqueries[:,0].cpu().unsqueeze(0)] )
imshow(torchvision.utils.make_grid(ct.view(2, nchannels,image_len,image_len)))<jupyter_output>Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).
{'fraction_non_alphanumeric': 0.11084885779313833, 'fraction_numerical': 0.012012138582146168, 'mean_word_length': 2.6469412849677223, 'pattern_counts': {'":': 5, '<': 24, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 43, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '192063', 'license': 'MIT', 'n_tokens_mistral': 7878, 'n_tokens_neox': 7496, 'n_words': 1547, 'path': 'LearnProjections-Beta.ipynb', 'repo_name': 'TrentBrick/attention-approximates-sdm', 'size': '23735'}