Custom Complex Loss Functions in Keras
Keras is a "building blocks" style deep learning framework that lets you build common deep learning models conveniently and intuitively. Even before TensorFlow came out, Keras was already almost the hottest deep learning framework at the time, using Theano as its backend, and nowadays Keras supports four backends simultaneously: Theano, TensorFlow, CNTK, and MXNet (the first three officially supported, MXNet not yet integrated officially) — which shows just how appealing Keras is.
Keras is indeed convenient, but that convenience doesn't come for free. One of the most commonly criticized shortcomings is its relatively low flexibility, making it hard to build complex models. It's true that Keras isn't particularly well suited to building complex models, but that's not to say it's impossible — it's just that the amount of code needed to build a very complex model in Keras ends up being not much less than writing it directly in TensorFlow. Still, Keras's friendly, convenient features (like that adorable training progress bar) mean there are always scenarios where we want to use it. So it's worth figuring out how to customize Keras models more flexibly. In this post we'll focus on custom loss functions.
Designing inputs and outputs
Keras models are functional — they have inputs, and they have outputs, and the loss is some error function between the predicted values and the true values. Keras itself comes with many built-in loss functions, such as MSE, cross-entropy, etc., which you can just call directly. To customize a loss, the most natural approach is to rewrite it by imitating Keras's built-in losses. more
For example, when doing classification, we often use a softmax output followed by cross-entropy as the loss. However, this approach has quite a few drawbacks, one of which is that the classification becomes overconfident — even with noisy input, the classification result ends up being almost either 1 or 0, which typically raises the risk of overfitting, and also makes it hard in practical applications to properly determine confidence intervals or set thresholds. So often we want to find ways to make the classifier less overconfident, and modifying the loss is one way to do this.
If we don't modify the loss, we're just using cross-entropy to fit a one-hot distribution. The formula for cross-entropy is
$$S(q|p)=-\sum_i q_i \log p_i$$
where $p_i$ is the predicted distribution and $q_i$ is the true distribution. For example, if the output is $[z_1,z_2,z_3]$ and the target is $[1,0,0]$, then
$$loss = -\log \Big(e^{z_1}/Z\Big),\, Z=e^{z_1}+e^{z_2}+e^{z_3}$$
As long as $z_1$ is already the maximum value of $[z_1,z_2,z_3]$, we can always "double down" — by increasing the number of training steps, we can make $z_1,z_2,z_3$ increase by a sufficiently large proportion (equivalently, increasing the norm of the vector $[z_1,z_2,z_3]$), so that $e^{z_1}/Z$ gets sufficiently close to 1 (equivalently, the loss gets sufficiently close to 0). This is where softmax's usual overconfidence comes from: as long as you blindly increase the norm, you can lower the loss, and the optimizer is only too happy to do so — the cost is far too low. To keep the classifier from becoming too confident, one solution is to not simply fit the one-hot distribution, but also spend a bit of effort fitting a uniform distribution as well, giving us a new loss:
$$loss = -(1-\varepsilon)\log \Big(e^{z_1}/Z\Big)-\varepsilon\sum_{i=1}^n \frac{1}{3}\log \Big(e^{z_i}/Z\Big),\, Z=e^{z_1}+e^{z_2}+e^{z_3}$$
This way, blindly increasing the ratio so that $e^{z_1}/Z$ approaches 1 is no longer the optimal solution, which can alleviate the overconfidence of softmax. In quite a few cases, this strategy can also increase test accuracy (by preventing overfitting).
So how should we write this in Keras? It's actually quite simple:
from keras.layers import Input,Embedding,LSTM,Dense
from keras.models import Model
from keras import backend as K
word_size = 128
nb_features = 10000
nb_classes = 10
encode_size = 64
input = Input(shape=(None,))
embedded = Embedding(nb_features,word_size)(input)
encoder = LSTM(encode_size)(embedded)
predict = Dense(nb_classes, activation='softmax')(encoder)
def mycrossentropy(y_true, y_pred, e=0.1):
loss1 = K.categorical_crossentropy(y_true, y_pred)
loss2 = K.categorical_crossentropy(K.ones_like(y_pred)/nb_classes, y_pred)
return (1-e)*loss1 + e*loss2
model = Model(inputs=input, outputs=predict)
model.compile(optimizer='adam', loss=mycrossentropy)
That is, we define a custom loss function that takes y_pred and y_true as inputs, and pass it into the model's compile step. Here, in mycrossentropy, the first term is the ordinary cross-entropy, and in the second term, we first construct a uniform distribution via K.ones_like(y_pred)/nb_classes, and then compute the cross-entropy between y_pred and this uniform distribution. That's all there is to it~
It's not just as simple as inputs and outputs
As mentioned above, Keras models have fixed inputs and outputs, and the loss is some error function between the predicted values and the true values. However, many models don't fit this pattern — for example, question-answering models with triplet loss.
Here we're talking about FAQ-style question answering with a fixed answer bank. A common approach for building such QA models is: first encode both the answers and the questions into vectors of the same length, then compare their cosine similarity — the larger the cosine, the better the match. This approach is easy to understand and is a fairly general framework — for instance, here the "questions" and "answers" don't necessarily need to be text; they could be images too, since only the encoding method differs, as long as we can ultimately encode something into a vector. But how do we train it? Naturally, we want the cosine value for the correct answer to be as large as possible, and the cosine value for wrong answers to be as small as possible, but this isn't strictly necessary. A more reasonable requirement is: the cosine value of the correct answer should be larger than the cosine values of all wrong answers — it doesn't matter by how much, even a tiny bit will do. This is exactly what leads to triplet loss:
$$loss = \max\Big(0, m+\cos(q,A_{\text{wrong}})-\cos(q,A_{\text{right}})\Big)$$
where $m$ is some positive number greater than zero.
How should we understand this loss? Note that we want to minimize the loss, so let's just look at the $m+\cos(q,A_{\text{wrong}})-\cos(q,A_{\text{right}})$ part: we know the goal is to widen the gap between the correct and wrong answers, but once $\cos(q,A_{\text{right}})-\cos(q,A_{\text{wrong}}) > m$ — that is, once the gap exceeds $m$ — because of the presence of $\max$, the loss becomes 0, at which point it has automatically reached its minimum and there's no more incentive to optimize it further. So the idea behind triplet loss is: we only want the correct answer to beat the wrong answer by a bit (not by as much as possible) — once we exceed $m$, we stop caring about it and instead focus our energy on the samples that haven't yet been separated!
Since we already have questions and correct answers, wrong answers can simply be picked at random, so constructing training samples this way is quite easy. But how do we implement triplet loss in Keras? It looks like a single-input, dual-output model, but it's not quite that simple — in Keras, a dual-output model can only have a separate loss assigned to each output, which are then summed with weights, but here the loss can't be expressed simply as a weighted sum of two terms. So how should we build such a model? Here's an example:
from keras.layers import Input,Embedding,LSTM,Dense,Lambda
from keras.layers.merge import dot
from keras.models import Model
from keras import backend as K
word_size = 128
nb_features = 10000
nb_classes = 10
encode_size = 64
margin = 0.1
embedding = Embedding(nb_features,word_size)
lstm_encoder = LSTM(encode_size)
def encode(input):
return lstm_encoder(embedding(input))
q_input = Input(shape=(None,))
a_right = Input(shape=(None,))
a_wrong = Input(shape=(None,))
q_encoded = encode(q_input)
a_right_encoded = encode(a_right)
a_wrong_encoded = encode(a_wrong)
q_encoded = Dense(encode_size)(q_encoded) #一般的做法是，直接讲问题和答案用同样的方法encode成向量后直接匹配，但我认为这是不合理的，我认为至少经过某个变换。
right_cos = dot([q_encoded,a_right_encoded], -1, normalize=True)
wrong_cos = dot([q_encoded,a_wrong_encoded], -1, normalize=True)
loss = Lambda(lambda x: K.relu(margin+x[0]-x[1]))([wrong_cos,right_cos])
model_train = Model(inputs=[q_input,a_right,a_wrong], outputs=loss)
model_q_encoder = Model(inputs=q_input, outputs=q_encoded)
model_a_encoder = Model(inputs=a_right, outputs=a_right_encoded)
model_train.compile(optimizer='adam', loss=lambda y_true,y_pred: y_pred)
model_q_encoder.compile(optimizer='adam', loss='mse')
model_a_encoder.compile(optimizer='adam', loss='mse')
model_train.fit([q,a1,a2], y, epochs=10)
#其中q,a1,a2分别是问题、正确答案、错误答案的batch，y是任意形状为(len(q),1)的矩阵
If you don't understand it the first time, please read it a few more times — this code embodies the general idea for implementing the most general kind of model in Keras: treat the target as an input, forming a multi-input model, and write the loss as a layer that becomes the final output. When building the model, you only need to define the model's output as the loss, and when compiling, simply set the loss to y_pred (since the model's output is the loss itself, y_pred is the loss), ignoring y_true entirely — during training, you can just feed in any array of the right shape for y_true. In the end, what we obtain are encoders for the question and the answer — that is, both questions and answers are each encoded into a vector, and we only need to compare $\cos$ to select the best answer.
The clever use of the Embedding layer
Before reading this section, please make sure you have a clear understanding of the Embedding layer. If you don't yet, please go read What Word Embedding and the Embedding Layer Really Are first. It bears repeating here: although word vectors are called "Word Embeddings," the Embedding layer is not the word vector, and has nothing whatsoever to do with word vectors!!! Don't ask silly questions like "how does this have anything to do with word vectors" — the Embedding layer has never had any direct connection with word vectors (it's just that it can be used when training word vectors). You can understand the Embedding layer in two equivalent ways: 1) it's an accelerated version of a fully-connected layer with one-hot input — that is, it's just a Dense layer whose input happens to be one-hot, and the two are mathematically completely equivalent; 2) it's simply a matrix lookup operation — you feed in an integer and it outputs the vector at the corresponding index, except that this matrix is trainable. (See — where's the connection to word vectors?)
In this section we'll look at center loss. As mentioned earlier, classification is usually done with softmax + cross-entropy; written in matrix form, softmax is
$$\text{softmax}\Big(\boldsymbol{W}\boldsymbol{x}+\boldsymbol{b}\Big)$$
where $\boldsymbol{x}$ can be understood as the extracted feature, and $\boldsymbol{W},\boldsymbol{b}$ is the weight of the final fully-connected layer, with the whole model trained jointly. The question is: what does the feature model $\boldsymbol{x}$ trained by this scheme actually look like?
In some cases, we care more about the feature $\boldsymbol{x}$ than about the final classification result. Take face recognition, for example: suppose we have a database of face images for 100,000 different people, with each person having several photos. Then we could train a 100,000-way classification model, and given a photo, determine which of the 100,000 people it belongs to. But this is only the training scenario — how do we apply it? In the actual deployment environment, say inside a company, there might only be a few hundred people; in a public security scenario, there might be several million people — so the 100,000-way classification model we trained is basically meaningless in either case. However, the features from just before the softmax in that model — that is, $\boldsymbol{x}$ mentioned in the previous paragraph — might still be quite meaningful. If, for the same person (i.e., the same class), $\boldsymbol{x}$ turns out to be basically the same, then in practical applications we can treat the trained model as a feature extraction tool, and simply run KNN (nearest neighbor) on the extracted features.
The idea sounds great, but reality is harsh: if you directly train with softmax, the resulting features don't necessarily have a clustering property — on the contrary, they tend to spread out to fill the entire space (leaving no room for other people; see the papers and articles related to center loss, e.g. this one). So how should we train in order to get results with a clustering property? Center loss uses a simple, crude, but effective scheme — adding a clustering penalty term. Written out in full, this is
$$loss = - \log\frac{e^{\boldsymbol{W}_y^{\top}\boldsymbol{x}+b_y}}{\sum\limits_i e^{\boldsymbol{W}_i^{\top}\boldsymbol{x}+b_i}} + \lambda \Big\Vert \boldsymbol{x}-\boldsymbol{c}_y \Big\Vert^2$$
where $y$ corresponds to the correct class. As you can see, the first term is just the ordinary softmax cross-entropy, and the second term is an extra penalty term: it defines a trainable center $\boldsymbol{c}$ for each class, and requires each class to stay close to its own center. So overall, the first term is responsible for pushing different classes apart, while the second term is responsible for pulling samples of the same class closer together.
So how do we implement this scheme in Keras? The key question is: how do we store the cluster centers? The answer is: the Embedding layer! As hinted at the start of this section, an Embedding is just a matrix waiting to be trained, which is exactly suited to storing the cluster center parameters. So, following the pattern from the second section, we get
from keras.layers import Input,Conv2D, MaxPooling2D,Flatten,Dense,Embedding,Lambda
from keras.models import Model
from keras import backend as K
nb_classes = 100
feature_size = 32
input_image = Input(shape=(224,224,3))
cnn = Conv2D(10, (2,2))(input_image)
cnn = MaxPooling2D((2,2))(cnn)
cnn = Flatten()(cnn)
feature = Dense(feature_size, activation='relu')(cnn)
predict = Dense(nb_classes, activation='softmax', name='softmax')(feature) #至此，得到一个常规的softmax分类模型
input_target = Input(shape=(1,))
centers = Embedding(nb_classes, feature_size)(input_target) #Embedding层用来存放中心
l2_loss = Lambda(lambda x: K.sum(K.square(x[0]-x[1][:,0]), 1, keepdims=True), name='l2_loss')([feature,centers])
model_train = Model(inputs=[input_image,input_target], outputs=[predict,l2_loss])
model_train.compile(optimizer='adam', loss=['sparse_categorical_crossentropy',lambda y_true,y_pred: y_pred], loss_weights=[1.,0.2], metrics={'softmax':'accuracy'})
model_predict = Model(inputs=input_image, outputs=predict)
model_predict.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model_train.fit([train_images,train_targets], [train_targets,random_y], epochs=10)
#TIPS：这里用的是sparse交叉熵，这样我们直接输入整数的类别编号作为目标，而不用转成one hot形式。所以Embedding层的输入，跟softmax的目标，都是train_targets，都是类别编号，而random_y是任意形状为(len(train_images),1)的矩阵。
Readers might wonder: why not write the overall loss as a single output and build the model that way, as we did with the triplet loss model in the second section, instead of making it dual-output like this?
As it happens, one important reason Keras enthusiasts are so fond of Keras is its progress bar — being able to display training loss and training accuracy in real time. If we wrote it the way we did in the second section, we wouldn't be able to set the metrics parameter, and so we wouldn't be able to display accuracy during training — which would be a bit of a shame. Writing it the way we do here, though, we can still see the training accuracy during training, and can also separately see the cross-entropy loss, the l2 loss, and the total loss — which is very satisfying.
Keras really is this much fun
With these three examples in hand, readers should now have a clear sense of the steps for building complex models in Keras — it should be said that this is actually fairly simple and flexible. Keras does have places where it's not flexible enough, but it's not nearly as incapable as some online comments suggest. On the whole, Keras is able to meet the needs of most people who want to quickly experiment with deep learning models. If you're still agonizing over which deep learning framework to choose, then just go with Keras — by the time you truly find that Keras can't meet your needs, you'll already be capable of handling any framework, and the dilemma will have resolved itself.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.