Tweaked training scripts

Author: david-baek

notebooks/modadd.ipynb

@@ -0,0 +1,8 @@
+## How to add new dataset for experiments
+
+1. Implement a function which returns the dataset dictionary in `utils/dataset.py`.
+2. Choose a unique id for the new dataset. Implement a function which evaluates the quality of representation in `utils/crystal_metric.py`. Modify the function `crystal_metric` to support the new data_id.
+3. Add the new data_id to the array `data_id_choices` in `run_exp.py`.
+4. If any auxiliary information is required to evaluate the representations, add them to the dictionary `aux_info` in `run_exp.py`. Sometimes, these information may depend on the specific dataset; In such cases, make any necessary modifications within each of the three experiment for loops in `run_exp.py`.
+5. Now, you're ready to test the new dataset! Command format is:
+`python run_exp.py --data_id new_data_id --model_id H_MLP`.

src/README.md

@@ -19,7 +19,7 @@
 model_id_choices = ["H_MLP", "standard_MLP", "H_transformer", "standard_transformer"]
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='Experiment')
-    parser.add_argument('--seed', type=int, default=77, help='random seed')
+    parser.add_argument('--seed', type=int, default=29, help='random seed')
     parser.add_argument('--data_id', type=str, required=True, choices=data_id_choices, help='Data ID')
     parser.add_argument('--model_id', type=str, required=True, choices=model_id_choices, help='Model ID')
 
@@ -42,6 +42,17 @@
     'embd_dim': 16,
 }
 
+aux_info = {}
+if data_id == "lattice":
+    aux_info["lattice_size"] = 5
+elif data_id == "greater":
+    aux_info["p"] = 30
+elif data_id == "equivalence":
+    aux_info["mod"] = 5
+elif data_id == "circle":
+    aux_info["p"] = 31
+else:
+    raise ValueError(f"Unknown data_id: {data_id}")
 
 # Train the model
 print(f"Training model with seed {seed}, data_id {data_id}, model_id {model_id}")
@@ -81,20 +92,9 @@
     torch.save(model.state_dict(), f"../results/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}.pt")
     with open(f"../results/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_train_results.json", "w") as f:
         json.dump(ret_dic["results"], f, indent=4)
-
-    aux_info = {}
-    if data_id == "lattice":
-        aux_info["lattice_size"] = 5
-    elif data_id == "greater":
-        aux_info["p"] = 30
-    elif data_id == "family_tree":
+    
+    if data_id == "family_tree":
         aux_info["dict_level"] = dataset['dict_level']
-    elif data_id == "equivalence":
-        aux_info["mod"] = 10
-    elif data_id == "circle":
-        aux_info["p"] = 59
-    else:
-        raise ValueError(f"Unknown data_id: {data_id}")
 
     if hasattr(model.embedding, 'weight'):
         metric_dict = crystal_metric(model.embedding.weight.cpu().detach(), data_id, aux_info)
@@ -128,19 +128,8 @@
     with open(f"../results/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_train_results.json", "w") as f:
         json.dump(ret_dic["results"], f, indent=4)
 
-    aux_info = {}
-    if data_id == "lattice":
-        aux_info["lattice_size"] = 5
-    elif data_id == "greater":
-        aux_info["p"] = 30
-    elif data_id == "family_tree":
+    if data_id == "family_tree":
         aux_info["dict_level"] = dataset['dict_level']
-    elif data_id == "equivalence":
-        aux_info["mod"] = 10
-    elif data_id == "circle":
-        aux_info["p"] = 59
-    else:
-        raise ValueError(f"Unknown data_id: {data_id}")
 
     if hasattr(model.embedding, 'weight'):
         metric_dict = crystal_metric(model.embedding.weight.cpu().detach(), data_id, aux_info)
@@ -179,19 +168,8 @@
     with open(f"../results/{seed}_{data_id}_{model_id}_{data_size}_{train_ratio}_train_results.json", "w") as f:
         json.dump(ret_dic["results"], f, indent=4)
 
-    aux_info = {}
-    if data_id == "lattice":
-        aux_info["lattice_size"] = 5
-    elif data_id == "greater":
-        aux_info["p"] = 30
-    elif data_id == "family_tree":
+    if data_id == "family_tree":
         aux_info["dict_level"] = dataset['dict_level']
-    elif data_id == "equivalence":
-        aux_info["mod"] = 10
-    elif data_id == "circle":
-        aux_info["p"] = 31
-    else:
-        raise ValueError(f"Unknown data_id: {data_id}")
 
     if hasattr(model.embedding, 'weight'):
         metric_dict = crystal_metric(model.embedding.weight.cpu().detach(), data_id, aux_info)

src/run_exp.py

@@ -19,7 +19,7 @@
 model_id_choices = ["H_MLP", "standard_MLP", "H_transformer", "standard_transformer"]
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='Experiment')
-    parser.add_argument('--seed', type=int, default=51, help='random seed')
+    parser.add_argument('--seed', type=int, default=29, help='random seed')
     parser.add_argument('--data_id', type=str, required=True, choices=data_id_choices, help='Data ID')
     parser.add_argument('--model_id', type=str, required=True, choices=model_id_choices, help='Model ID')
 
@@ -65,7 +65,7 @@
 if data_id == "lattice":
     aux_info["lattice_size"] = 5
 elif data_id == "greater":
-    aux_info["p"] = 200
+    aux_info["p"] = 30
 elif data_id == "family_tree":
     aux_info["dict_level"] = dataset['dict_level']
 elif data_id == "equivalence":

src/unit_exp.py

@@ -99,45 +99,39 @@ def greater_metric(reps, aux_info):
 
 def family_tree_metric(reps, aux_info):
 
+    dict_level = aux_info['dict_level']
+    reps = reps[:(max(dict_level.keys()) + 1)]
+
     pca = PCA(n_components=min(reps.shape[0], reps.shape[1]))
     reps = pca.fit_transform(reps)
     reps = reps[:, :2]
 
-    dict_level = aux_info['dict_level']
-
-    # Group individuals by generation
-    generation_groups = {}
-    for individual, generation in dict_level.items():
-        if generation not in generation_groups:
-            generation_groups[generation] = []
-        generation_groups[generation].append(individual)
 
-
-    # Compute the collinearity of representations for individuals within the same generation
-    collinearity_by_generation = {}
-
-    for generation, individuals in generation_groups.items():
-        # Get the indices of individuals in this generation
-        indices = [individual for individual in individuals]
-        # Extract their representations
-        gen_representations = reps[indices]
-
-        # Compute collinearity by fixing one vector as a pivot
-        if gen_representations.shape[0] > 2:  # Ensure there are at least three individuals
-            pivot = gen_representations[1] - gen_representations[0]  # Difference between first two vectors
-            dot_products = (gen_representations[2:] - gen_representations[0]) @ pivot
-            norms = np.linalg.norm((gen_representations[2:] - gen_representations[0]), axis=1) * np.linalg.norm(pivot)
-            
-            norms = np.where(norms == 0, np.nan, norms)
-            collinearity = np.abs(dot_products / norms)  # Cosine similarity with the pivot
-            collinearity = np.nan_to_num(collinearity, nan=1.0)
-            collinearity_by_generation[generation] = collinearity.mean()
+    # Group embeddings by generation
+    levels = {}
+    for node, generation in dict_level.items():
+        if generation not in levels:
+            levels[generation] = []
+        levels[generation].append(reps[node])
+    
+    # Compute one-dimensionality for each generation
+    level_scores = {}
+    for generation, points in levels.items():
+        if len(points) < 5:
+            continue
+        
+        points_array = np.stack(points)  # Convert to NumPy array
+        pca_sub = PCA(n_components=min(points_array.shape[0], points_array.shape[1]))
+        pca_sub.fit(points_array)
+        one_dimensionality = pca_sub.explained_variance_ratio_[0]  # Ratio of variance explained by the first PC
+        level_scores[generation] = one_dimensionality
 
 
+#    pca.fit_transform(reps)
     variances = pca.explained_variance_ratio_
 
     metric_dict = {
-        'metric': float(1 - np.mean([collinearity for collinearity in collinearity_by_generation.values() if not np.isnan(collinearity)])),
+        'metric': float(1 - np.mean(list(level_scores.values()))),
         'variances': variances.tolist(),
     }
     return metric_dict
@@ -156,6 +150,10 @@ def equivalence_metric(reps, aux_info):
             else:
                 diff_arr.append(np.linalg.norm(reps[i] - reps[j]))
 
+    # Filter Outliers
+    diff_arr = np.array(diff_arr)
+    diff_arr = diff_arr[diff_arr < np.mean(cross_diff_arr)]
+
     pca = PCA(n_components=min(reps.shape[0], reps.shape[1]))
     emb_pca = pca.fit_transform(reps)
     variances = pca.explained_variance_ratio_
@@ -174,18 +172,28 @@ def circle_metric(reps, aux_info):
     emb_pca = pca.fit_transform(reps)
     variances = pca.explained_variance_ratio_
 
+    points = emb_pca[:, :2]
+
+    min_x, min_y = points.min(axis=0)
+    max_x, max_y = points.max(axis=0)
+    width = max_x - min_x
+    height = max_y - min_y
+    
+    # Normalize points to [0, 1] in both dimensions
+    normalized_points = (points - [min_x, min_y]) / [width, height]
+
     # Compute the centroid of the points
-    centroid = np.mean(emb_pca, axis=0)
+    centroid = np.mean(normalized_points, axis=0)
 
     # Compute distances of points from the centroid
-    distances = np.linalg.norm(emb_pca - centroid, axis=1)
+    distances = np.linalg.norm(normalized_points - centroid, axis=1)
 
     # Mean and standard deviation of distances
     mean_distance = np.mean(distances)
     std_distance = np.std(distances)
 
     # Circularity score
-    circularity_score = 1 - (std_distance / mean_distance)
+    circularity_score = (std_distance / mean_distance)
 
 
     metric_dict = {

src/utils/crystal_metric.py

@@ -151,7 +151,7 @@ def is_desc(a, b):
     data_id = torch.from_numpy(x).to(device)
     labels = torch.from_numpy(target).to(device)
 
-    vocab_size = p+2
+    vocab_size = p
 
     dataset = {}
     dataset['data_id'] = data_id

src/utils/dataset.py

@@ -63,15 +63,14 @@ def train_single_model(param_dict: dict):
         dataset = parallelogram_dataset(p=5, dim=2, num=data_size, seed=seed, device=device)
         input_token = 3
     elif data_id == "greater":
-        dataset = greater_than_dataset(p=200, num=data_size, seed=seed, device=device)
+        dataset = greater_than_dataset(p=30, num=data_size, seed=seed, device=device)
     elif data_id == "family_tree":
         dataset = family_tree_dataset_2(p=127, num=data_size, seed=seed, device=device)
     elif data_id == "equivalence":
         input_token = 2
-        dataset = mod_classification_dataset(p=300, num=data_size, seed=seed, device=device)
+        dataset = mod_classification_dataset(p=100, num=data_size, seed=seed, device=device)
     elif data_id == "circle":
         dataset = modular_addition_dataset(p=31, num=data_size, seed=seed, device=device)
-        input_token = 3
     else:
         raise ValueError(f"Unknown data_id: {data_id}")
 
@@ -92,9 +91,9 @@ def train_single_model(param_dict: dict):
         shp = [input_token * embd_dim, hidden_size, embd_dim, vocab_size]
         model = MLP(shp=shp, vocab_size=vocab_size, embd_dim=embd_dim, input_token=input_token, unembd=unembd, weight_tied=weight_tied, seed=seed).to(device)
     elif model_id == "H_transformer":
-        model = ToyTransformer(vocab_size=vocab_size, d_model=embd_dim, nhead=16, num_layers=3, seq_len=input_token, use_dist_layer=True).to(device)
+        model = ToyTransformer(vocab_size=vocab_size, d_model=embd_dim, nhead=8, num_layers=1, seq_len=input_token, use_dist_layer=True).to(device)
     elif model_id == "standard_transformer":
-        model = ToyTransformer(vocab_size=vocab_size, d_model=embd_dim, nhead=16, num_layers=3, seq_len=input_token, use_dist_layer=False).to(device)
+        model = ToyTransformer(vocab_size=vocab_size, d_model=embd_dim, nhead=8, num_layers=1, seq_len=input_token, use_dist_layer=False).to(device)
     else:
         raise ValueError(f"Unknown model_id: {model_id}")
 
@@ -106,7 +105,7 @@ def train_single_model(param_dict: dict):
     test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
 
     ret_dic = {}
-    ret_dic["results"] = model.train(param_dict={'num_epochs': 10000, 'learning_rate': 0.01, 'train_dataloader': train_dataloader, 'test_dataloader': test_dataloader, 'device': device})
+    ret_dic["results"] = model.train(param_dict={'num_epochs': 4000, 'learning_rate': 0.001, 'train_dataloader': train_dataloader, 'test_dataloader': test_dataloader, 'device': device})
     ret_dic["model"] = model
     ret_dic["dataset"] = dataset
 

src/utils/driver.py

@@ -29,7 +29,7 @@ def train(self, param_dict: dict):
         test_accuracies = []       
 
         best_loss = float('inf')
-        patience = 200
+        patience = 100
         min_delta = 1e-4
         counter = 0 
 
@@ -57,7 +57,7 @@ def train(self, param_dict: dict):
                 else:
                     total_loss = loss + lamb_reg * torch.mean(torch.sqrt(torch.mean(self.embedding.data**2, dim=0)))
 
-                loss.backward()
+                total_loss.backward()
                 optimizer.step()
                 train_loss += loss.item()
 
@@ -100,9 +100,11 @@ def train(self, param_dict: dict):
             else:
                 counter += 1  # Increment counter if no improvement
 
+            '''
             if counter >= patience:
                 print("Early stopping triggered!")
                 break
+            '''
 
         ret_dic = {}
         ret_dic['train_losses'] = train_losses