Customizing Complex Loss Functions in Keras

Input-Output Design

Keras models are functional, meaning they have inputs and outputs, and the loss is a certain error function between the predicted values and the true values. Keras itself comes with many loss functions, such as MSE and cross-entropy, which can be called directly. To customize a loss, the most natural method is to rewrite it following the pattern of Keras’s built-in losses.

For example, when dealing with classification problems, we often use softmax output and cross-entropy as the loss. However, this approach has several drawbacks. One is that the classification is "too confident"; even with noisy input, the classification results are almost always either 0 or 1. This usually leads to a risk of overfitting and makes it difficult to determine confidence intervals or set thresholds in practical applications. Therefore, we often want to find ways to make the classification less confident, and modifying the loss is one such method.

If we do not modify the loss, we use 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 among [z_1, z_2, z_3], we can always "exacerbate the situation"—by increasing the training steps to make z_1, z_2, z_3 increase by a sufficiently large proportion (equivalently, increasing the magnitude of the vector [z_1, z_2, z_3]), so that e^{z_1}/Z is close enough to 1 (equivalently, the loss is close enough to 0). This is the source of why softmax is usually overconfident: as long as the magnitude is blindly increased, the loss can be reduced. The trainer is certainly happy to do this as the cost is too low. To prevent the classification from being too confident, one solution is not to fit the one-hot distribution purely, but to spend some effort fitting a uniform distribution. The new loss becomes: 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} In this way, blindly increasing the proportion to make e^{z_1}/Z approach 1 is no longer the optimal solution, thereby alleviating the overconfidence of softmax. In many cases, this strategy can also increase test accuracy (preventing overfitting).

So, how do we write this in Keras? It is 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)

This involves defining a custom loss function with inputs y_true and y_pred and passing it into the model’s compile method. In mycrossentropy, the first term is the ordinary cross-entropy, and in the second term, a uniform distribution is constructed via K.ones_like(y_pred)/nb_classes, and then the cross-entropy between y_pred and the uniform distribution is calculated. It’s that simple!

More Than Just Input and Output

As mentioned earlier, Keras models have fixed inputs and outputs, and the loss is an error function between predictions and ground truth. However, many models do not follow this pattern, such as Question-Answering (QA) models and Triplet Loss.

This problem refers to FAQ-style QA with a fixed answer database. A common method for building QA models is: first, encode both the answer and the question into vectors of the same length, then compare their cosine values; the larger the cosine, the better the match. This approach is easy to understand and is a general framework; for instance, the questions and answers do not necessarily have to be text—images work too, as long as the encoding method is appropriate and a vector can eventually be produced. But how do we train it? We naturally want the cosine value of the correct answer to be as large as possible and the cosine value of incorrect answers to be as small as possible, but this is not strictly necessary. A reasonable requirement is: the cosine value of the correct answer should be larger than that of all incorrect answers; the magnitude of the difference doesn’t matter, even a tiny bit is enough. This leads to the Triplet Loss: loss = \max\Big(0, m+\cos(q,A_{\text{wrong}})-\cos(q,A_{\text{right}})\Big) where m is a positive constant greater than zero.

How do we understand this loss? Since we want to minimize the loss, we only look at the part m+\cos(q,A_{\text{wrong}})-\cos(q,A_{\text{right}}). We know the goal is to widen the gap between the correct and incorrect answers. However, once \cos(q,A_{\text{right}})-\cos(q,A_{\text{wrong}}) > m, i.e., the gap is greater than m, the loss becomes 0 due to the \max function. At this point, the minimum is automatically reached, and it will no longer be optimized. Therefore, the idea of Triplet Loss is: we only want the gap between the correct and incorrect answers to be slightly larger (not necessarily as large as possible). Once it exceeds m, stop worrying about it and focus on samples that haven’t been separated yet!

Since we already have questions and correct answers, incorrect answers can simply be picked at random, making it very easy to construct training samples. But how do we implement Triplet Loss in Keras? It looks like a single-input, dual-output model, but it’s not that simple. In Keras, a dual-output model can only have separate losses set for each output, which are then weighted and summed. Here, it cannot be simply expressed as a weighted sum of two terms. How should we build such a model? Here is 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)

# Common practice is to match vectors directly after encoding, 
# but I believe a transformation is often necessary.
q_encoded = Dense(encode_size)(q_encoded) 

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')

# q, a1, a2 are batches of questions, correct answers, and wrong answers.
# y is an arbitrary matrix of shape (len(q), 1).
model_train.fit([q,a1,a2], y, epochs=10)

If you don’t understand it at first, please read it several times. This code contains the logic for implementing the most general models in Keras: Treat the target as an input to form a multi-input model, write the loss as a layer as the final output, and when building the model, simply define the model’s output as the loss. During compilation, set the loss directly to y_pred (since the model’s output is the loss, y_pred is the loss), ignoring y_true. During training, just pass any array of the correct shape for y_true. Finally, we obtain encoders for questions and answers. By encoding them into vectors and comparing their cosines, we can select the optimal answer.

Clever Use of Embedding Layers

Before reading this section, please ensure you have a clear understanding of the Embedding layer. If not, please refer to "What exactly are Word Vectors and Embeddings?". It must be emphasized repeatedly that although word vectors are called Word Embeddings, the Embedding layer is not a word vector and has nothing to do with them! Do not ask silly questions like "how did it get related to 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 ways: 1. It is an accelerated version of a fully connected layer with one-hot input; that is, it is a Dense layer with one-hot input, mathematically equivalent. 2. It is a matrix lookup operation: input an integer, output the vector at the corresponding index, except this matrix is trainable. (See, where is the connection to word vectors?)

In this part, we focus on Center Loss. As mentioned, classification is usually done with softmax + cross-entropy. In matrix form, softmax is: \text{softmax}\Big(\boldsymbol{W}\boldsymbol{x}+\boldsymbol{b}\Big) where \boldsymbol{x} can be understood as the extracted features, and \boldsymbol{W}, \boldsymbol{b} are the weights of the final fully connected layer. The entire model is trained together. The question is: what form do the features \boldsymbol{x} trained by such a scheme take?

In some cases, we care more about the features \boldsymbol{x} than the final classification result. For example, in face recognition, suppose we have a database of 100,000 different people, each with several photos. We could train a 100,000-class model to judge which of the 100,000 people a given photo belongs to. But this is only the training scenario. How do we apply it? In a specific application environment, such as inside a company, there might be only a few hundred people; in a public security detection scenario, there might be millions. Thus, the 100,000-class model is basically meaningless, but the features before the softmax, \boldsymbol{x}, might still be very meaningful. If for the same person (the same class), \boldsymbol{x} is basically the same, then in practice, we can use the trained model as a feature extraction tool and use the extracted features directly with KNN (K-Nearest Neighbors).

The vision is beautiful, but reality is harsh. Training softmax directly does not necessarily result in features with clustering properties; instead, they tend to fill the entire space (leaving no room for others; refer to Center Loss papers or articles). So, how do we train to achieve clustering properties? Center Loss uses a simple but effective solution—adding a clustering penalty term. Written completely, it 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 we can see, the first term is the ordinary softmax cross-entropy, and the second term is an additional penalty. It defines trainable centers \boldsymbol{c} for each class, requiring each class to be very close to its respective center. Thus, overall, the first term is responsible for increasing the distance between different classes, while the second term is responsible for reducing the distance within the same class.

How do we implement this in Keras? The key is: where to store the clustering centers? The answer is the Embedding layer! As hinted at the beginning of this section, an Embedding is just a matrix to be trained, which is perfect for storing clustering center parameters. Thus, mimicking the approach in 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)
# Conventional softmax classification model
predict = Dense(nb_classes, activation='softmax', name='softmax')(feature) 

input_target = Input(shape=(1,))
# Embedding layer used to store centers
centers = Embedding(nb_classes, feature_size)(input_target) 
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'])

# TIPS: sparse cross-entropy allows integer class labels as targets.
# train_targets are used for both softmax and Embedding input.
# random_y is an arbitrary matrix of shape (len(train_images), 1).
model_train.fit([train_images,train_targets], [train_targets,random_y], epochs=10)

Readers might wonder why we didn’t write the overall loss as a single output like in the Triplet Loss model, but instead used a dual-output approach.

In fact, Keras enthusiasts love Keras largely because of its progress bar—it can display training loss and accuracy in real-time. If written as in the second section, the metrics parameter cannot be set, and training accuracy cannot be displayed during the process, which would be a small regret. By writing it this way, we can still see the training accuracy and separately see the cross-entropy loss, l2_loss, and total loss, which is very comfortable.

Keras is Just This Fun

With these three cases, readers should have a good idea of the steps for building complex models in Keras. It should be said that it is relatively simple and flexible. Keras does have areas where it is not flexible enough, but it is not as "incapable" as some online comments suggest. Overall, Keras is capable of meeting the needs of most people for rapid experimentation with deep learning models. If you are still hesitating about the choice of deep learning frameworks, choose Keras—by the time you truly feel Keras cannot meet your needs, you will already have the ability to master any framework, and the hesitation will be gone.

When reprinting, please include the original address of this article: https://kexue.fm/archives/4493

For more detailed reprinting matters, please refer to: Scientific Space FAQ