Single input - multiple outputs with different loss functions in Keras: how is the gradient computed?

Question

I've implemented a neural network with single input - multiple outputs using Keras API. The general structure of the network is like in this figure:

Because each branch does a different task, I choose different loss functions (cross-entropy for the classifier and MSE for the regressor). The code is quite lengthy, but I can summarize it as follows:

# define the input layer
inputs = Input(shape=(..))

# define the common part (from the input layer)
common = make_common_part(inputs)

# model with a single input until the common part, then branches out
model = Model(inputs=inputs, outputs=[make_branch_1(common), make_branch_2(common)])
 
# compile the model with MSE for 
model.compile(loss=['categorical_crossentropy', 'mean_squared_error'], optimizer='adam')

The code works and returns the results as expected, but I want to understand how Keras works in the background.

My intuition is that this model can learn the features (in the common part) that are good for both the classification and regression tasks because its weight matrix $W$ is updated by the gradients of the losses from the 2 branches. I guess the gradient used for backpropagation is computed on the sum of two losses:

$$ L_{total} = L_{crossentropy} + L_{mse} = \sum_{i=1}^{N} CE(y_{reg}^{(i)}, \hat{y}_{reg}^{(i)}) + \sum_{i=1}^{N} MSE(y_{clf}^{(i)}, \hat{y}_{clf}^{(i)}) $$

I tried to confirm this by reading the code of keras and tensorflow but it's still unclear to me. It seems to me the losses are computed separately but I don't understand how the gradients in form of tensors are computed in tensorflow... Are the gradient of both losses summed up or are they left separated?

Any insights you could give me are very much appreciated! Thank you very much in advance!

Answer 1

In the code you provided, Keras is using a multi-output architecture for your neural network, with two branches each having their own output and loss function. The common part of the network is shared between the two branches, and the gradients from the two branches are backpropagated separately to update the weights in the common part of the network.

The gradients from the two branches are computed separately using the respective loss functions, and are then used to update the weights in the common part of the network. Keras uses the optimizer specified in the compile function to update the weights. In this case, the optimizer is 'adam'.

In terms of the gradients computation, each branch has its own gradients computed and they are not summed up. The gradients are used to update the weights on the common part of the network, the gradients are computed based on the respective losses.

The gradients are passed to the optimizer which updates the weights by minimizing the total loss. The optimizer is responsible for updating the weights based on the gradients it receives. Therefore, the optimizer will update the weights based on the gradients computed by the two branches separately.

So in summary, the losses are computed separately for each branch, and the gradients from each branch are used to update the weights in the common part of the network separately, but the optimizer takes the total loss which is the sum of the two losses and uses the gradients to optimize the model.