How does one design a custom loss function? What features make a loss function "good"?

Question

I have a custom situation for which I am trying to design a cost function.

The idea is that you have a stack of LSTMs doing something slightly unconventional. Each LSTM$_l$ computes a linear transformation of its hidden layer $V_{l-1}h^t_l$ to predict the hidden layer vector $h^{t+1}_{l-1}$ at the level below; at the same time, it should not stray far from the value it was predicted to have from the layer above (i.e., we want $h^t_l \approx V_lh^{t-1}_{l+1}$). That is, the goal at each layer is that:

$V_{l-1}h^t_l \approx h^{t+1}_{l-1}$

and

$h^t_l \approx V_lh^{t-1}_{l+1}$

I would like to design a cost function that will bring about this situation in a system of $L$ stacked LSTMs. I introduced an error variable at each layer $E^{t}_{l} = V_lh^{t-1}_{l+1} - h^{t}_{l}$, where $l$ ranges from $0$ to $L-1$, such that $h^t_0 \equiv x^t$, and the goal is to model a sequence $x^{1:T}$. The loss function is then:

$\Large \mathcal{L}_{\text{subtract}} = \sum_{t=1}^T\sum_{i=0}^{L-1}|E^t_i|$

And the sequence of steps is the following:

However, this simple subtractive loss function is not very good, in the sense that the model could not even overfit a minibatch (I did this as a test). I tried a divisive loss:

$\Large E^t_l = \frac{V_{l}h_{l+1}^{t-1}}{\max(\epsilon, h_{l}^t)} + \frac{h_{l}^t}{\max(\epsilon,V_{l}h_{l+1}^{t-1})}$

But the model also did not learn anything (probably because division led to numerical instability). How do I go about designing a stable loss function that I can use at all layers?

NB My goal for this project is to build a language model based on predictive coding principles, to test a hypothesis about brain activity (i.e., the goal is not to beat SoTA).

Answer 1

To answer the titular question, a key characteristic of a loss function is that the loss is minimized at the target values $y$. In other words, if you're estimating a quantity, the least loss should be assigned to the estimates that are exactly correct.

Using $\mathcal{L}_\text{subtract}$ and divisive losses are not loss functions because they do not involve the target variables. In other words, you can achieve the absolute minimum value of 0 without learning anything about what you want to model.

Consider this loss function $L(\hat{y}) = (\hat{y} - c)^2$ for model outputs $\hat{y}$, a constant $c$ and model targets $y$. This function is clearly convex in $\hat{y},c$, but it has nothing to do with the model targets $y$. If $c$ is a fixed target and our model outputs $\hat{y}$, then we achieve a minimum at $\hat{y}=c$, and this is true no matter what the targets $y$ happen to be.

My suggestion is to use $\mathcal{L}_\text{subtract}$ as a regularization to augment a typical loss function. So if you're using a binomial cross entropy loss, you could write down a total loss as

$$ \text{total loss} = \text{BCE}(y,\hat{y}) + \lambda \mathcal{L}_\text{subtract}$$

for $\lambda >0$ some tuning parameter controlling how much deviation can occur from one layer to the next. Likewise, you can use an alternative loss in place of $\text{BCE}$ to model other types of data.

---

This seems related to, but not exactly the same as "Regularizing RNNs by Stabilizing Activations" by David Krueger & Roland Memisevic. Perhaps their work and bibliography is of interest.