Pythorch: Dictionary learning with Neural Networks. Adam Opt does not converge

Question

I have implemented a code based on the paper of https://arxiv.org/pdf/1707.00225.pdf

However, the algorithm 1 of the paper (which is basically 2 steps: first step you calculate K with the parameters of the NN, and second step is to apply gradient descent to find the loss) works for a simple model with discrete spectrum. But, when I apply it to a more complex case (Duffing equation), It does not converge. I was wondering if you have any suggestions to make it work.

import gc
import os
import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim


from matplotlib import pyplot as plt
from scipy.stats import uniform
from statistics import mean
from torch.optim.lr_scheduler import StepLR
from torch.optim.optimizer import Optimizer




data_name = 'Duffing_oscillator' 



lambda_regula = 0.01  # Regularization and Tikhonov regulatization

epsilon = 30

d = 2
l = 170
M = 22
N = 10000
I = torch.eye(M + 3, M + 3)
learning_rate = 0.001


model = nn.Sequential(
    nn.Linear(d, l),
    nn.Tanh(),
    nn.Linear(l, l),
    nn.Tanh(),
    nn.Linear(l, l),
    nn.Tanh(),
    nn.Linear(l, M),
)





optimizer = optim.AdamW(model.parameters(), lr=learning_rate) 




loss_fn = torch.nn.MSELoss(reduction='sum')

def data_generation(training_data, cut):
       data = np.loadtxt(('./data/%s_%s.csv' % (data_name, training_data)), delimiter=',', dtype=np.float64)[:cut]
       data = torch.tensor(data, dtype=torch.float32)
       return data

def Frobenius_norm(X):
    M = torch.mm(X, torch.transpose(X, 0, 1))
    return torch.sum(torch.diag(M, 0))


   


x_data =data_generation("train_x", N)
y_data =data_generation("train_y", N)
    

    
#Fixed dictionary:  one constant map and two coordinate projection maps  

fixed_sai = torch.tensor([i + [0.5] for i in x_data.detach().tolist()], dtype=torch.float32)
y_fixed_sai = torch.tensor([i + [0.5] for i in y_data.detach().tolist()], dtype=torch.float32)



loss = float("INF")

# for i in range(iteration):
while loss > epsilon:

    pred_sai = model(x_data.detach())  
    y_pred_sai = model(y_data.detach())

   

    pred_sai = torch.cat([pred_sai, fixed_sai], dim=1)
    y_pred_sai = torch.cat([y_pred_sai, y_fixed_sai], dim=1)

    pred_sai_T = torch.transpose(pred_sai, 0, 1)
    y_pred_sai_T = torch.transpose(y_pred_sai, 0, 1)
 
    #Step 1   
    
    G=torch.linalg.pinv(1/N  * torch.mm(pred_sai_T, pred_sai)  + lambda_regula * I)
    A= 1/N  * torch.mm(pred_sai_T, y_pred_sai)

    K_tilde = torch.mm(G, A)
    K_tilde = K_tilde.clone().detach()



    # loss calculation for the algorithm
    MSE = (y_pred_sai_T - torch.mm(K_tilde, pred_sai_T))** 2
    loss = torch.sum(MSE) + lambda_regula * Frobenius_norm(K_tilde) ** 2
    

    #Step 2   

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    print(loss)    
                          

I have tried changing optimizers, varying the learning rate ...

Answer 1

Have you tried changing your input data?

I think that the next part of the code should be after the while statement

optimizer.zero_grad()