4-neural-network.py

import os
import time
import argparse
import numpy as np
import torch
from tqdm import tqdm
from pathlib import Path

from dataset import StateDataset
from constants import DEFAULT_LLM_MODEL_PATH
from utils import *

MAX_EPOCHS = 50
DEFAULT_LR = 1e-5


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        'datasets',
        nargs='+',
    )
    parser.add_argument('-g', '--gpu', type=int, required=True)
    parser.add_argument('-l', '--llm_path', type=str, default=DEFAULT_LLM_MODEL_PATH)
    parser.add_argument('-k', '--nb_fs_examples', type=int, default=0)
    parser.add_argument('-s', '--seed', type=int, default=42)
    parser.add_argument('--lr', type=float, default=DEFAULT_LR)
    parser.add_argument('--normalize', action='store_true')
    args = parser.parse_args()
    # args = parser.parse_args(args=['aqua', '-g', '0'])
    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
    llm_name = Path(args.llm_path).stem
    sample_folder = Path(f"samples/{llm_name}/{args.nb_fs_examples}")
    results_folder = Path(f"results/{llm_name}/{args.nb_fs_examples}")
    
    for dataset_name in args.datasets:
        # ----------------------- 1. Select dataset ----------------------- 
        # Load results dict
        results = load_results_dict(dataset_name, folder=results_folder)
        with open(f"datasets/{dataset_name}/split.pkl", "rb") as file:
            I_train, I_valid, I_test = pickle.load(file)
        
        # -------- AQuA --------
        if dataset_name=="aqua":
            train_data = StateDataset(sample_folder/"aqua/train", I_train)
            valid_data = StateDataset(sample_folder/"aqua/train", I_valid)
            test_data = StateDataset(sample_folder/"aqua/train", I_test)
        # -------- CosmosQA --------
        elif dataset_name=="cosmosqa":
            train_data = StateDataset(sample_folder/"cosmosqa/train", I_train)
            valid_data = StateDataset(sample_folder/"cosmosqa/train", I_valid)
            test_data = StateDataset(sample_folder/"cosmosqa/valid")
        # -------- MMLU --------
        elif dataset_name=="mmlu":
            train_data = StateDataset(sample_folder/"mmlu/train")
            valid_data = StateDataset(sample_folder/"mmlu/valid")
            test_data = StateDataset(sample_folder/"mmlu/test")
        # -------- MedMCQA --------
        elif dataset_name=="medmcqa":
            train_data = StateDataset(sample_folder/"medmcqa/train", I_train)
            valid_data = StateDataset(sample_folder/"medmcqa/train", I_valid)
            test_data = StateDataset(sample_folder/"medmcqa/train", I_test)
        # -------- HellaSwag --------
        elif dataset_name=="hellaswag":
            train_data = StateDataset(sample_folder/"hellaswag/train", I_train)
            valid_data = StateDataset(sample_folder/"hellaswag/train", I_valid)
            test_data = StateDataset(sample_folder/"hellaswag/valid")
        # -------- LogiQA --------
        elif dataset_name=="logiqa":
            train_data = StateDataset(sample_folder/"logiqa/train")
            valid_data = StateDataset(sample_folder/"logiqa/valid")
            test_data = StateDataset(sample_folder/"logiqa/test")
        # -------- Unrecognized --------
        else:
            raise ValueError(f"Unknown dataset '{dataset_name}'")

        # ----------------------- 2. Load data in memory ----------------------- 
        
        print(f"Loading dataset '{dataset_name}' in memory")
        state, logits, answer = train_data[0] # To know the shape for allocating the tensors
        NB_LAYERS, STATE_SIZE = state.shape
        NB_CHOICES = logits.shape[-1]

        time_start = time.perf_counter()
        nb_train = len(train_data)
        train_states = torch.empty(nb_train, *state.shape, dtype=torch.float16)
        train_logits = torch.empty(nb_train, *logits.shape, dtype=torch.float16)
        train_answers = torch.empty(nb_train, dtype=int)
        for i in tqdm(range(nb_train), dynamic_ncols=True):
            state, logits, answer = train_data[i]
            train_states[i] = torch.from_numpy(state)
            train_logits[i] = torch.from_numpy(logits)
            train_answers[i] = answer
        train_data = torch.utils.data.TensorDataset(train_states, train_logits, train_answers)
        results["time"]["train_data_loading"] = time.perf_counter() - time_start

        time_start = time.perf_counter()
        nb_valid = len(valid_data)
        valid_states = torch.empty(nb_valid, *state.shape, dtype=torch.float16)
        valid_logits = torch.empty(nb_valid, *logits.shape, dtype=torch.float16)
        valid_answers = torch.empty(nb_valid, dtype=int)
        for i in tqdm(range(nb_valid), dynamic_ncols=True):
            state, logits, answer = valid_data[i]
            valid_states[i] = torch.from_numpy(state)
            valid_logits[i] = torch.from_numpy(logits)
            valid_answers[i] = answer
        valid_data = torch.utils.data.TensorDataset(valid_states, valid_logits, valid_answers)
        results["time"]["valid_data_loading"] = time.perf_counter() - time_start

        nb_test = len(test_data)
        test_states = torch.empty(nb_test, *state.shape, dtype=torch.float16)
        test_logits = torch.empty(nb_test, *logits.shape, dtype=torch.float16)
        test_answers = torch.empty(nb_test, dtype=int)
        for i in tqdm(range(nb_test), dynamic_ncols=True):
            state, logits, answer = test_data[i]
            test_states[i] = torch.from_numpy(state)
            test_logits[i] = torch.from_numpy(logits)
            test_answers[i] = answer
        test_data = torch.utils.data.TensorDataset(test_states, test_logits, test_answers)

        train_dataloader = torch.utils.data.DataLoader(train_data, shuffle=True, batch_size=32, pin_memory=True)
        valid_dataloader = torch.utils.data.DataLoader(valid_data, batch_size=128, pin_memory=True)
        test_dataloader = torch.utils.data.DataLoader(test_data, batch_size=128,pin_memory=True)
        
        if args.normalize:
            states_center = train_states.mean(0, keepdim=True)
            states_scale = train_states.std(0, keepdim=True)
            states_scale[states_scale.abs()<1e-6] = 1
            logits_center = train_logits.mean(0, keepdim=True)
            logits_scale = train_logits.std(0, keepdim=True)
            logits_scale[logits_scale.abs()<1e-6] = 1
        else:
            states_center = torch.zeros(1, *train_states.shape[1:])
            states_scale = torch.ones(1, *train_states.shape[1:])
            logits_center = torch.zeros(1, *train_logits.shape[1:])
            logits_scale = torch.ones(1, *train_logits.shape[1:])

        # ----------------------- 3.Train neural networks -----------------------
        
        results["info"]["nn_lr"] = args.lr
        results["info"]["nn_normalize"] = args.normalize
        
        print(f"\n---------------------- NN on logits ----------------------")
        np.random.seed(args.seed)
        torch.manual_seed(args.seed)
        nn_model = torch.nn.Sequential(
                    torch.nn.LayerNorm(NB_CHOICES),
                    torch.nn.Linear(NB_CHOICES, 32),
                    torch.nn.ReLU(),
                    torch.nn.Linear(32, NB_CHOICES)
                ).to("cuda")
        nn_optimizer = torch.optim.Adam(nn_model.parameters(), lr=args.lr)

        time_start = time.perf_counter()
        train_dists, train_accuracies, train_nces = [], [], []
        valid_dists, valid_accuracies, valid_nces = [], [], []
        test_dists, test_accuracies, test_nces = [], [], []
        best_valid_accuracy = 0
        best_valid_nce = np.inf
        for epoch in range(MAX_EPOCHS):
            dist, accuracy, nce = [], [], []
            for batch_id, (_, batch_logits, batch_answers) in enumerate(train_dataloader):
                # if batch_id*32 > 10000:
                #     break
                batch_logits = batch_logits-logits_center
                batch_logits = (batch_logits/logits_scale).float().to("cuda")
                batch_answers = batch_answers.to("cuda")
                batch_predicted_logits = nn_model(batch_logits)
                batch_dist = batch_predicted_logits.softmax(dim=-1)
                batch_accuracy = (batch_predicted_logits.argmax(-1) == batch_answers)
                batch_nce = -batch_dist.gather(dim=-1, index=batch_answers.unsqueeze(-1)).squeeze(-1).log()

                nn_optimizer.zero_grad()
                batch_nce.mean().backward()
                nn_optimizer.step()

                dist += batch_dist.detach().cpu()
                accuracy += batch_accuracy.detach().cpu()
                nce += batch_nce.detach().cpu()
            train_dists += [torch.stack(dist)]
            train_accuracies += [torch.stack(accuracy)]
            train_nces += [torch.stack(nce)]

            dist, accuracy, nce = [], [], []
            for _, batch_logits, batch_answers in valid_dataloader:
                with torch.no_grad():
                    batch_logits = batch_logits-logits_center
                    batch_logits = (batch_logits/logits_scale).float().to("cuda")
                    batch_answers = batch_answers.to("cuda")
                    batch_predicted_logits = nn_model(batch_logits)
                    batch_dist = batch_predicted_logits.softmax(dim=-1)
                    batch_accuracy = (batch_predicted_logits.argmax(-1) == batch_answers)
                    batch_nce = -batch_dist.gather(dim=-1, index=batch_answers.unsqueeze(-1)).squeeze(-1).log()

                dist += batch_dist.detach().cpu()
                accuracy += batch_accuracy.detach().cpu()
                nce += batch_nce.detach().cpu()
            valid_dists += [torch.stack(dist)]
            valid_accuracies += [torch.stack(accuracy)]
            valid_nces += [torch.stack(nce)]

            dist, accuracy, nce = [], [], []
            for _, batch_logits, batch_answers in test_dataloader:
                with torch.no_grad():
                    batch_logits = batch_logits-logits_center
                    batch_logits = (batch_logits/logits_scale).float().to("cuda")
                    batch_answers = batch_answers.to("cuda")
                    batch_predicted_logits = nn_model(batch_logits)
                    batch_dist = batch_predicted_logits.softmax(dim=-1)
                    batch_accuracy = (batch_predicted_logits.argmax(-1) == batch_answers)
                    batch_nce = -batch_dist.gather(dim=-1, index=batch_answers.unsqueeze(-1)).squeeze(-1).log()

                dist += batch_dist.detach().cpu()
                accuracy += batch_accuracy.detach().cpu()
                nce += batch_nce.detach().cpu()
            test_dists += [torch.stack(dist)]
            test_accuracies += [torch.stack(accuracy)]
            test_nces += [torch.stack(nce)]

            new_record = False
            if valid_accuracies[-1].float().mean(0) >= best_valid_accuracy:
                new_record = True
                best_valid_nce = valid_nces[-1].mean(0)
                best_valid_accuracy = valid_accuracies[-1].float().mean(0)

                final_test_dists = test_dists[-1]
                final_test_accuracies = test_accuracies[-1]
                final_test_nces = test_nces[-1]
                final_nb_epochs = epoch

            print(f"Epoch {epoch}"
                  f" | Train NegLogProb {train_nces[-1].mean(0):.4f}, Accuracy {train_accuracies[-1].float().mean(0):.4f}"
                  f" | Valid NegLogProb {valid_nces[-1].mean(0):.4f}, Accuracy {valid_accuracies[-1].float().mean(0):.4f}"
                  f" | Test NegLogProb {test_nces[-1].mean(0):.4f}, Accuracy {test_accuracies[-1].float().mean(0):.4f}"
                  f"{" [*]" if new_record else ""}")

        print(f"Final test NegLogProb {final_test_nces.mean(0):.4f}, Accuracy: {final_test_accuracies.float().mean(0):.4f}")

        method_name="logits_nn"
        results["time"][method_name] = time.perf_counter() - time_start
        results[method_name]["accuracies"] = final_test_accuracies
        results[method_name]["distributions"] = final_test_dists
        results[method_name]["neglogprob"] = final_test_nces
        results[method_name]["nb_epochs"] = final_nb_epochs
        save_results_dict(results, dataset_name, folder=results_folder)


        print(f"\n---------------------- NN on last state ----------------------")
        np.random.seed(args.seed)
        torch.manual_seed(args.seed)
        nn_model = torch.nn.Sequential(
                    torch.nn.LayerNorm(STATE_SIZE),
                    torch.nn.Linear(STATE_SIZE, 32),
                    torch.nn.ReLU(),
                    torch.nn.Linear(32, NB_CHOICES)
                ).to("cuda")
        nn_optimizer = torch.optim.Adam(nn_model.parameters(), lr=args.lr)

        time_start = time.perf_counter()
        train_dists, train_accuracies, train_nces = [], [], []
        valid_dists, valid_accuracies, valid_nces = [], [], []
        test_dists, test_accuracies, test_nces = [], [], []
        best_valid_accuracy = 0
        best_valid_nce = np.inf
        for epoch in range(MAX_EPOCHS):
            dist, accuracy, nce = [], [], []
            for batch_id, (batch_states, _, batch_answers) in enumerate(train_dataloader):
                batch_states = batch_states[:, -1, :]-states_center[:, -1, :]
                batch_states = (batch_states/states_scale[:, -1, :]).float().to("cuda")
                batch_answers = batch_answers.to("cuda")
                batch_predicted_logits = nn_model(batch_states)
                batch_dist = batch_predicted_logits.softmax(dim=-1)
                batch_accuracy = (batch_predicted_logits.argmax(-1) == batch_answers)
                batch_nce = -batch_dist.gather(dim=-1, index=batch_answers.unsqueeze(-1)).squeeze(-1).log()

                nn_optimizer.zero_grad()
                batch_nce.mean().backward()
                nn_optimizer.step()

                dist += batch_dist.detach().cpu()
                accuracy += batch_accuracy.detach().cpu()
                nce += batch_nce.detach().cpu()
            train_dists += [torch.stack(dist)]
            train_accuracies += [torch.stack(accuracy)]
            train_nces += [torch.stack(nce)]

            dist, accuracy, nce = [], [], []
            for batch_states, _, batch_answers in valid_dataloader:
                with torch.no_grad():
                    batch_states = batch_states[:, -1, :]-states_center[:, -1, :]
                    batch_states = (batch_states/states_scale[:, -1, :]).float().to("cuda")
                    batch_answers = batch_answers.to("cuda")
                    batch_predicted_logits = nn_model(batch_states)
                    batch_dist = batch_predicted_logits.softmax(dim=-1)
                    batch_accuracy = (batch_predicted_logits.argmax(-1) == batch_answers)
                    batch_nce = -batch_dist.gather(dim=-1, index=batch_answers.unsqueeze(-1)).squeeze(-1).log()

                dist += batch_dist.detach().cpu()
                accuracy += batch_accuracy.detach().cpu()
                nce += batch_nce.detach().cpu()
            valid_dists += [torch.stack(dist)]
            valid_accuracies += [torch.stack(accuracy)]
            valid_nces += [torch.stack(nce)]

            dist, accuracy, nce = [], [], []
            for batch_states, _, batch_answers in test_dataloader:
                with torch.no_grad():
                    batch_states = batch_states[:, -1, :]-states_center[:, -1, :]
                    batch_states = (batch_states/states_scale[:, -1, :]).float().to("cuda")
                    batch_answers = batch_answers.to("cuda")
                    batch_predicted_logits = nn_model(batch_states)
                    batch_dist = batch_predicted_logits.softmax(dim=-1)
                    batch_accuracy = (batch_predicted_logits.argmax(-1) == batch_answers)
                    batch_nce = -batch_dist.gather(dim=-1, index=batch_answers.unsqueeze(-1)).squeeze(-1).log()

                dist += batch_dist.detach().cpu()
                accuracy += batch_accuracy.detach().cpu()
                nce += batch_nce.detach().cpu()
            test_dists += [torch.stack(dist)]
            test_accuracies += [torch.stack(accuracy)]
            test_nces += [torch.stack(nce)]

            new_record = False
            if valid_accuracies[-1].float().mean(0) >= best_valid_accuracy:
                new_record = True
                best_valid_nce = valid_nces[-1].mean(0)
                best_valid_accuracy = valid_accuracies[-1].float().mean(0)
                best_weights = nn_model.state_dict()

                final_test_dists = test_dists[-1]
                final_test_accuracies = test_accuracies[-1]
                final_test_nces = test_nces[-1]
                final_nb_epochs = epoch

            print(f"Epoch {epoch}"
                  f" | Train NegLogProb {train_nces[-1].mean(0):.4f}, Accuracy {train_accuracies[-1].float().mean(0):.4f}"
                  f" | Valid NegLogProb {valid_nces[-1].mean(0):.4f}, Accuracy {valid_accuracies[-1].float().mean(0):.4f}"
                  f" | Test NegLogProb {test_nces[-1].mean(0):.4f}, Accuracy {test_accuracies[-1].float().mean(0):.4f}"
                  f"{" [*]" if new_record else ""}")

        nn_model.load_state_dict(best_weights)
        print(f"Final test NegLogProb {final_test_nces.mean(0):.4f}, Accuracy: {final_test_accuracies.float().mean(0):.4f}")

        method_name="last_nn"
        results["time"][method_name] = time.perf_counter() - time_start
        results[method_name]["accuracies"] = final_test_accuracies
        results[method_name]["distributions"] = final_test_dists
        results[method_name]["neglogprob"] = final_test_nces
        results[method_name]["nb_epochs"] = final_nb_epochs
        
        # Compute gradients on test set, for analysis
        grads = []
        for batch_states, _, batch_answers in test_dataloader:
            batch_states = batch_states[:, -1, :]-states_center[:, -1, :]
            batch_states = (batch_states/states_scale[:, -1, :]).float().to("cuda")
            batch_states.requires_grad = True
            batch_answers = batch_answers.to("cuda")
            batch_predicted_logits = nn_model(batch_states)

            batch_dist = batch_predicted_logits.softmax(dim=-1)
            batch_nce = -batch_dist.gather(dim=-1, index=batch_answers.unsqueeze(-1)).squeeze(-1).log()
            batch_grads = torch.autograd.grad(inputs=batch_states, outputs=batch_nce.mean())[0].norm(2, dim=-1)
            grads += batch_grads.detach().cpu()
        test_grads = torch.stack(grads)
        test_center = test_states[:, -1, :].mean(0)
        test_scale = test_states[:, -1, :].std(0)
        test_scale[test_scale.abs()<1e-6] = 1
        results["analysis"][method_name+"_grads"] = test_grads
        results["analysis"][method_name+"_scale"] = test_scale
        
        # Save results
        save_results_dict(results, dataset_name, folder=results_folder)

        print(f"\n---------------------- NN on last 10 states ----------------------")
        np.random.seed(args.seed)
        torch.manual_seed(args.seed)
        class NNModel(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.layer1 = torch.nn.Sequential(
                    torch.nn.LayerNorm(STATE_SIZE),
                    torch.nn.Linear(STATE_SIZE, 32),
                    torch.nn.ReLU()
                )
                self.layer2 = torch.nn.Sequential(
                    torch.nn.LayerNorm(10),
                    torch.nn.Linear(10, 8),
                    torch.nn.ReLU()
                )
                self.layer3 = torch.nn.Sequential(
                    torch.nn.LayerNorm(8*32),
                    torch.nn.Linear(8*32, NB_CHOICES)
                )

            def forward(self, input_):
                hidden = self.layer1(input_)
                hidden = self.layer2(hidden.transpose(-1, -2))
                hidden = hidden.reshape(*hidden.shape[:-2], 8*32)
                output = self.layer3(hidden)
                return output
        nn_model = NNModel().to("cuda")
        nn_optimizer = torch.optim.Adam(nn_model.parameters(), lr=args.lr)

        time_start = time.perf_counter()
        train_dists, train_accuracies, train_nces = [], [], []
        valid_dists, valid_accuracies, valid_nces = [], [], []
        test_dists, test_accuracies, test_nces = [], [], []
        best_valid_accuracy = 0
        best_valid_nce = np.inf
        for epoch in range(MAX_EPOCHS):
            dist, accuracy, nce = [], [], []
            for batch_id, (batch_states, _, batch_answers) in enumerate(train_dataloader):
                batch_states = batch_states[:, -10:, :]-states_center[:, -10:, :]
                batch_states = (batch_states/states_scale[:, -10:, :]).float().to("cuda")
                batch_answers = batch_answers.to("cuda")
                batch_predicted_logits = nn_model(batch_states)
                batch_dist = batch_predicted_logits.softmax(dim=-1)
                batch_accuracy = (batch_predicted_logits.argmax(-1) == batch_answers)
                batch_nce = -batch_dist.gather(dim=-1, index=batch_answers.unsqueeze(-1)).squeeze(-1).log()

                nn_optimizer.zero_grad()
                batch_nce.mean().backward()
                nn_optimizer.step()

                dist += batch_dist.detach().cpu()
                accuracy += batch_accuracy.detach().cpu()
                nce += batch_nce.detach().cpu()
            train_dists += [torch.stack(dist)]
            train_accuracies += [torch.stack(accuracy)]
            train_nces += [torch.stack(nce)]

            dist, accuracy, nce = [], [], []
            for batch_states, _, batch_answers in valid_dataloader:
                with torch.no_grad():
                    batch_states = batch_states[:, -10:, :]-states_center[:, -10:, :]
                    batch_states = (batch_states/states_scale[:, -10:, :]).float().to("cuda")
                    batch_answers = batch_answers.to("cuda")
                    batch_predicted_logits = nn_model(batch_states)
                    batch_dist = batch_predicted_logits.softmax(dim=-1)
                    batch_accuracy = (batch_predicted_logits.argmax(-1) == batch_answers)
                    batch_nce = -batch_dist.gather(dim=-1, index=batch_answers.unsqueeze(-1)).squeeze(-1).log()

                dist += batch_dist.detach().cpu()
                accuracy += batch_accuracy.detach().cpu()
                nce += batch_nce.detach().cpu()
            valid_dists += [torch.stack(dist)]
            valid_accuracies += [torch.stack(accuracy)]
            valid_nces += [torch.stack(nce)]

            dist, accuracy, nce = [], [], []
            for batch_states, _, batch_answers in test_dataloader:
                with torch.no_grad():
                    batch_states = batch_states[:, -10:, :]-states_center[:, -10:, :]
                    batch_states = (batch_states/states_scale[:, -10:, :]).float().to("cuda")
                    batch_answers = batch_answers.to("cuda")
                    batch_predicted_logits = nn_model(batch_states)
                    batch_dist = batch_predicted_logits.softmax(dim=-1)
                    batch_accuracy = (batch_predicted_logits.argmax(-1) == batch_answers)
                    batch_nce = -batch_dist.gather(dim=-1, index=batch_answers.unsqueeze(-1)).squeeze(-1).log()

                dist += batch_dist.detach().cpu()
                accuracy += batch_accuracy.detach().cpu()
                nce += batch_nce.detach().cpu()
            test_dists += [torch.stack(dist)]
            test_accuracies += [torch.stack(accuracy)]
            test_nces += [torch.stack(nce)]

            new_record = False
            if valid_accuracies[-1].float().mean(0) >= best_valid_accuracy:
                new_record = True
                best_valid_nce = valid_nces[-1].mean(0)
                best_valid_accuracy = valid_accuracies[-1].float().mean(0)
                best_weights = nn_model.state_dict()

                final_test_dists = test_dists[-1]
                final_test_accuracies = test_accuracies[-1]
                final_test_nces = test_nces[-1]
                final_nb_epochs = epoch

            print(f"Epoch {epoch}"
                  f" | Train NegLogProb {train_nces[-1].mean(0):.4f}, Accuracy {train_accuracies[-1].float().mean(0):.4f}"
                  f" | Valid NegLogProb {valid_nces[-1].mean(0):.4f}, Accuracy {valid_accuracies[-1].float().mean(0):.4f}"
                  f" | Test NegLogProb {test_nces[-1].mean(0):.4f}, Accuracy {test_accuracies[-1].float().mean(0):.4f}"
                  f"{" [*]" if new_record else ""}")

        nn_model.load_state_dict(best_weights)
        print(f"Final test NegLogProb {final_test_nces.mean(0):.4f}, Accuracy: {final_test_accuracies.float().mean(0):.4f}")

        method_name="last10_nn"
        results["time"][method_name] = time.perf_counter() - time_start
        results[method_name]["accuracies"] = final_test_accuracies
        results[method_name]["distributions"] = final_test_dists
        results[method_name]["neglogprob"] = final_test_nces
        results[method_name]["nb_epochs"] = final_nb_epochs
        
        # Compute gradients on validation set, for analysis
        grads = []
        for batch_states, _, batch_answers in test_dataloader:
            batch_states = batch_states[:, -10:, :]-states_center[:, -10:, :]
            batch_states = (batch_states/states_scale[:, -10:, :]).float().to("cuda")
            batch_states.requires_grad = True
            batch_answers = batch_answers.to("cuda")
            batch_predicted_logits = nn_model(batch_states)

            batch_dist = batch_predicted_logits.softmax(dim=-1)
            batch_nce = -batch_dist.gather(dim=-1, index=batch_answers.unsqueeze(-1)).squeeze(-1).log()
            batch_grads = torch.autograd.grad(inputs=batch_states, outputs=batch_nce.mean())[0].norm(2, dim=-1)
            grads += batch_grads.detach().cpu()
        test_grads = torch.stack(grads)
        test_center = test_states[:, -10:, :].mean(0)
        test_scale = test_states[:, -10:, :].std(0)
        test_scale[test_scale.abs()<1e-6] = 1
        results["analysis"][method_name+"_grads"] = test_grads
        results["analysis"][method_name+"_scale"] = test_scale
        
        # Save results
        save_results_dict(results, dataset_name, folder=results_folder)

        
        print(f"\n---------------------- NN on all states ----------------------")
        np.random.seed(args.seed)
        torch.manual_seed(args.seed)
        class NNModel(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.layer1 = torch.nn.Sequential(
                    torch.nn.LayerNorm(STATE_SIZE),
                    torch.nn.Linear(STATE_SIZE, 32),
                    torch.nn.ReLU()
                )
                self.layer2 = torch.nn.Sequential(
                    torch.nn.LayerNorm(NB_LAYERS),
                    torch.nn.Linear(NB_LAYERS, 8),
                    torch.nn.ReLU()
                )
                self.layer3 = torch.nn.Sequential(
                    torch.nn.LayerNorm(8*32),
                    torch.nn.Linear(8*32, NB_CHOICES)
                )

            def forward(self, input_):
                hidden = self.layer1(input_)
                hidden = self.layer2(hidden.transpose(-1, -2))
                hidden = hidden.reshape(*hidden.shape[:-2], 8*32)
                output = self.layer3(hidden)
                return output
        nn_model = NNModel().to("cuda")
        nn_optimizer = torch.optim.Adam(nn_model.parameters(), lr=args.lr)

        time_start = time.perf_counter()
        train_dists, train_accuracies, train_nces = [], [], []
        valid_dists, valid_accuracies, valid_nces = [], [], []
        test_dists, test_accuracies, test_nces = [], [], []
        best_valid_accuracy = 0
        best_valid_nce = np.inf
        for epoch in range(MAX_EPOCHS):
            dist, accuracy, nce = [], [], []
            for batch_id, (batch_states, _, batch_answers) in enumerate(train_dataloader):
                batch_states = batch_states-states_center
                batch_states = (batch_states/states_scale).float().to("cuda")
                batch_answers = batch_answers.to("cuda")
                batch_predicted_logits = nn_model(batch_states)
                batch_dist = batch_predicted_logits.softmax(dim=-1)
                batch_accuracy = (batch_predicted_logits.argmax(-1) == batch_answers)
                batch_nce = -batch_dist.gather(dim=-1, index=batch_answers.unsqueeze(-1)).squeeze(-1).log()

                nn_optimizer.zero_grad()
                batch_nce.mean().backward()
                nn_optimizer.step()

                dist += batch_dist.detach().cpu()
                accuracy += batch_accuracy.detach().cpu()
                nce += batch_nce.detach().cpu()
            train_dists += [torch.stack(dist)]
            train_accuracies += [torch.stack(accuracy)]
            train_nces += [torch.stack(nce)]

            dist, accuracy, nce = [], [], []
            for batch_states, _, batch_answers in valid_dataloader:
                with torch.no_grad():
                    batch_states = batch_states-states_center
                    batch_states = (batch_states/states_scale).float().to("cuda")
                    batch_answers = batch_answers.to("cuda")
                    batch_predicted_logits = nn_model(batch_states)
                    batch_dist = batch_predicted_logits.softmax(dim=-1)
                    batch_accuracy = (batch_predicted_logits.argmax(-1) == batch_answers)
                    batch_nce = -batch_dist.gather(dim=-1, index=batch_answers.unsqueeze(-1)).squeeze(-1).log()

                dist += batch_dist.detach().cpu()
                accuracy += batch_accuracy.detach().cpu()
                nce += batch_nce.detach().cpu()
            valid_dists += [torch.stack(dist)]
            valid_accuracies += [torch.stack(accuracy)]
            valid_nces += [torch.stack(nce)]

            dist, accuracy, nce = [], [], []
            for batch_states, _, batch_answers in test_dataloader:
                with torch.no_grad():
                    batch_states = batch_states-states_center
                    batch_states = (batch_states/states_scale).float().to("cuda")
                    batch_answers = batch_answers.to("cuda")
                    batch_predicted_logits = nn_model(batch_states)
                    batch_dist = batch_predicted_logits.softmax(dim=-1)
                    batch_accuracy = (batch_predicted_logits.argmax(-1) == batch_answers)
                    batch_nce = -batch_dist.gather(dim=-1, index=batch_answers.unsqueeze(-1)).squeeze(-1).log()

                dist += batch_dist.detach().cpu()
                accuracy += batch_accuracy.detach().cpu()
                nce += batch_nce.detach().cpu()
            test_dists += [torch.stack(dist)]
            test_accuracies += [torch.stack(accuracy)]
            test_nces += [torch.stack(nce)]

            new_record = False
            if valid_accuracies[-1].float().mean(0) >= best_valid_accuracy:
                new_record = True
                best_valid_nce = valid_nces[-1].mean(0)
                best_valid_accuracy = valid_accuracies[-1].float().mean(0)
                best_weights = nn_model.state_dict()

                final_test_dists = test_dists[-1]
                final_test_accuracies = test_accuracies[-1]
                final_test_nces = test_nces[-1]
                final_nb_epochs = epoch

            print(f"Epoch {epoch}"
                  f" | Train NegLogProb {train_nces[-1].mean(0):.4f}, Accuracy {train_accuracies[-1].float().mean(0):.4f}"
                  f" | Valid NegLogProb {valid_nces[-1].mean(0):.4f}, Accuracy {valid_accuracies[-1].float().mean(0):.4f}"
                  f" | Test NegLogProb {test_nces[-1].mean(0):.4f}, Accuracy {test_accuracies[-1].float().mean(0):.4f}"
                  f"{" [*]" if new_record else ""}")

        nn_model.load_state_dict(best_weights)
        print(f"Final test NegLogProb {final_test_nces.mean(0):.4f}, Accuracy: {final_test_accuracies.float().mean(0):.4f}")

        method_name="all_nn"
        results["time"][method_name] = time.perf_counter() - time_start
        results[method_name]["accuracies"] = final_test_accuracies
        results[method_name]["distributions"] = final_test_dists
        results[method_name]["neglogprob"] = final_test_nces
        results[method_name]["nb_epochs"] = final_nb_epochs
        
        # Compute gradients on test set, for analysis
        grads = []
        for batch_states, _, batch_answers in test_dataloader:
            batch_states = batch_states-states_center
            batch_states = (batch_states/states_scale).float().to("cuda")
            batch_states.requires_grad = True
            batch_answers = batch_answers.to("cuda")
            batch_predicted_logits = nn_model(batch_states)

            batch_dist = batch_predicted_logits.softmax(dim=-1)
            batch_nce = -batch_dist.gather(dim=-1, index=batch_answers.unsqueeze(-1)).squeeze(-1).log()
            batch_grads = torch.autograd.grad(inputs=batch_states, outputs=batch_nce.mean())[0].norm(2, dim=-1)
            grads += batch_grads.detach().cpu()
        test_grads = torch.stack(grads)
        test_center = test_states.mean(0)
        test_scale = test_states.std(0)
        test_scale[test_scale.abs()<1e-6] = 1
        results["analysis"][method_name+"_grads"] = test_grads
        results["analysis"][method_name+"_scale"] = test_scale
        
        # Save results
        save_results_dict(results, dataset_name, folder=results_folder)