What should I think about when designing a custom loss function?

Question

I'm trying to get my toy network to learn a sine wave.

I output (via tanh) a number between -1 and 1, and I want the network to minimise the following loss, where self(x) are the predictions.

loss = -torch.mean(self(x)*y)

This should be equivalent to trading a stock with a sinusoidal price.

The issue I'm having is that the network doesn't learn anything. It does work if I change the loss function to be torch.mean((self(x)-y)**2) (MSE), but this isn't what I want. I'm trying to focus the network on 'making a profit', not making a prediction.

I think the issue may be related to the convexity of the loss function, but I'm not sure, and I'm not certain how to proceed. I've experimented with differing learning rates, but alas nothing works.

What should I be thinking about?

Actual code:

%load_ext tensorboard
import matplotlib.pyplot as plt; plt.rcParams["figure.figsize"] = (30,8)
import torch;from torch.utils.data import Dataset, DataLoader
import torch.nn.functional as F;import pytorch_lightning as pl
from torch import nn, tensor
def piecewise(x): return 2*(x>0)-1

class TsDs(torch.utils.data.Dataset):
  def __init__(self, s, l=5): super().__init__();self.l,self.s=l,s
  def __len__(self): return self.s.shape[0] - 1 - self.l
  def __getitem__(self, i): return self.s[i:i+self.l], torch.log(self.s[i+self.l+1]/self.s[i+self.l])
  def plt(self): plt.plot(self.s)

class TsDm(pl.LightningDataModule):
  def __init__(self, length=5000, batch_size=1000): super().__init__();self.batch_size=batch_size;self.s = torch.sin(torch.arange(length)*0.2) + 5 + 0*torch.rand(length)
  def train_dataloader(self): return DataLoader(TsDs(self.s[:3999]), batch_size=self.batch_size, shuffle=True)
  def val_dataloader(self): return DataLoader(TsDs(self.s[4000:]), batch_size=self.batch_size)

dm = TsDm()

class MyModel(pl.LightningModule):
    def __init__(self, learning_rate=0.01):
        super().__init__();self.learning_rate = learning_rate
        super().__init__();self.learning_rate = learning_rate
        self.conv1 = nn.Conv1d(1,5,2)
        self.lin1 = nn.Linear(20,3);self.lin2 = nn.Linear(3,1)
        # self.network = nn.Sequential(nn.Conv1d(1,5,2),nn.ReLU(),nn.Linear(20,3),nn.ReLU(),nn.Linear(3,1), nn.Tanh())
        # self.network = nn.Sequential(nn.Linear(5,5),nn.ReLU(),nn.Linear(5,3),nn.ReLU(),nn.Linear(3,1), nn.Tanh())
    def forward(self, x): 
        out = x.unsqueeze(1)
        out = self.conv1(out)
        out = out.reshape(-1,20)
        out = nn.ReLU()(out)
        out = self.lin1(out)
        out = nn.ReLU()(out)
        out = self.lin2(out)
        return nn.Tanh()(out)

    def step(self, batch, batch_idx, stage):
        x, y = batch
        loss = -torch.mean(self(x)*y)
        # loss = torch.mean((self(x)-y)**2)
        print(loss)
        self.log("loss", loss, prog_bar=True)
        return loss
    def training_step(self, batch, batch_idx): return self.step(batch, batch_idx, "train")
    def validation_step(self, batch, batch_idx): return self.step(batch, batch_idx, "val")
    def configure_optimizers(self): return torch.optim.SGD(self.parameters(), lr=self.learning_rate)

#logger = pl.loggers.TensorBoardLogger(save_dir="/content/")
mm = MyModel(0.1);trainer = pl.Trainer(max_epochs=10)
# trainer.tune(mm, dm)
trainer.fit(mm, datamodule=dm)
# 

Answer 1

It doesn't matter that your loss is not convex. As a matter of fact, the loss function of a neural network is in general neither convex nor concave (reference).

As ImotVoksim points out, the issue is that the loss function you've defined has nothing to do with the problem you're trying to solve.

For example, a stock price of zero is going to give you a loss of zero and hence no gradients at all: the neural network is allowed to output arbitrary values whenever the stock price is zero.

You want to "make a profit". I'm not sure why the MSE is not good in this case: if your neural network outputs the correct price of the stock for the next time period, you can use this information to make the trade which will maximize your profit.

Or do you want to predict the price further in the future? In that case you could use an MSE of the form

torch.mean((self(x_t)-y_{t+n})**2)

where x_t is the input at time period t and y_{t+n} the price of the stock at time period t+n, n a number you choose.

Answer 2

The loss function you have defined is the negative mean of the predictions multiplied by the targets, where both values are on the closed interval [-1,1]. The alternative you've listed is the MSE. Let's look at what these two loss functions result in with some concrete values:

self(x)	y	Your Loss	MSE
1	-1	1	4
0.5	-1	0.5	2.25
0	-1	0	1
-0.5	-1	-0.5	0.25
-1	-1	-1	0

A good loss function is such that when we are making good predictions the loss is close to 0, and when we are making bad predictions the loss increases above 0. That's why we often call the lost function a cost function. In other words, bad predictions are costly for the model, good predictions cost nothing ($L = 0$). As you can see, your loss function lacks this property, and instead it attains 0 whenever your prediction or the target is 0. Hence why your model cannot learn.