Multiple parameters recovery using Deep Learning
As a simplified version of my actual research problem, let's say I have a second-order polynomial function y = ax^2 + bx + c and I want to use a deep neural network to predict the parameters a, b and c given the variable x and the value of the function y. The variable x and the parameters a,b,c are exctracted from a uniform distribution in the range [0,1].
When I try to train the network using different architectures, cost functions and hyperparameters combinations among the most used, I always got the same issue: the train and test losses rapidly converge to a value significantly higher than 0, then starts to fluctuate in a strange way and the predictions are not accurate (see figures as a general example, the predictions for b are similar, c is slightly better but still not satisfactory). This happens even if I set higher momentum or lower learning rates. Also, I got the same issue if I try to recover one parameter at a time.

As an example, here is the PyTorch code I used for my first test (4 layers, first 3 followed by ReLU, MSELoss, RMSprop optimizer with learning rate = 0.001 and momentum 0.9).
class PRNet(nn.Module):
    def __init__(self, input_size, output_size):
        super(PRNet, self).__init__()
        self.input_size = input_size
        self.fc1   = nn.Linear(self.input_size, 32)
        self.relu1 = nn.ReLU()
        self.fc2   = nn.Linear(32, 64)
        self.relu2 = nn.ReLU()
        self.fc3   = nn.Linear(64, 64)
        self.relu3 = nn.ReLU()
        self.fc4   = nn.Linear(64, output_size)

    def forward(self, x):
        output = self.fc1(x)
        output = self.relu1(output)
        output = self.fc2(output)
        output = self.relu2(output)
        output = self.fc3(output)
        output = self.relu3(output)
        output = self.fc4(output)
        return output

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

var_x    = np.random.rand(100000)
pars_abc = np.random.rand(3, 100000)
func_y   = pars[0]*var**2 + pars[1] * var + pars[2]

data = np.vstack((var_x, func_y)).T
parameters = pars_abc.T

X = torch.Tensor(data).to(device).float()
y = torch.Tensor(parameters).to(device).float()

train_size       = int(0.8 * len(data))
batch_size       = 100
train_dataset    = TensorDataset(X[:train_size], y[:train_size])
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=False)

prnet = PRNet(X.shape[1], 3).to(device)

loss_function = nn.MSELoss()
optimizer = torch.optim.RMSprop(prnet.parameters(), lr=1e-4, momentum=0.9)

num_epochs = 25

for epoch in range(0, num_epochs):
    print(f'Starting epoch {epoch+1}')
    current_loss = 0.0
    
    for i, batch in enumerate(train_dataloader, 0):
        inputs, targets = batch

        optimizer.zero_grad()

        outputs = prnet(inputs)
        test_outputs = prnet(X[train_size:].to(device))
       
        train_loss = loss_function(outputs, targets)
        test_loss  = loss_function(test_outputs, y[train_size:])

        train_loss_plot[epoch,i] = train_loss.item()
        test_loss_plot[epoch,i]  = test_loss.item()

        train_loss.backward()
        optimizer.step()

What could be the cause of this issue? Are the features not representative enough? Do I need a custom loss more suitable for this problem?

During training, when a model's loss starts fluctuating, the most probable cause for such a pattern to show up is that the learning rate is high for the weights to get to the required value.
Consider this example. Suppose in your model, a parameter (weight), initialized with a value of 0.1, needs to get to a value of 0.00423 and the learning rate is set to 0.001.
Now, let's assume that the parameter has reached a value of 0.004 after a few epochs of training. Gradient descent will try to increase the value in order to make it equal to the target value but since the learning rate is only upto 3 decimal digits, the parameter value will now become 0.005. Since the value has now increased, gradient descent will try to decrease the value which will change the parameter value back to 0.004 and thus starting a fluctuation pattern.
To solve this issue, using a small learning rate will not help. Because if you use a small learning rate then the model will learn too slowly and might not converge at all. What you are probably looking for is a way to use a variable learning rate policy in your training. With such a policy, you can begin with a large learning rate initially so that the model learns faster. And later on, when the model parameters get close to the target values, the learning rate should decrease automatically in order to make the parameters reach as close as possible to the target. These policies are called learning rate schedulers.
There are several functions in PyTorch that let you use a learning rate scheduler of your choice. You can look for them in their documentation.
I'll suggest you to go for the Reduce LR on Plateau scheduler. It will let you set a threshold and a factor. Whenever your model loss does not improve over the specified threshold of number of epochs, it will decrease the learning rate by the factor.
https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.ReduceLROnPlateau.html#torch.optim.lr_scheduler.ReduceLROnPlateau
