Matrix Tensor multiplication in YOTO implementation

I am trying to implement Loss conditional training (YOTO) which allows you to train a whole spectrum of losses by using a loss function that takes parameters. I am training an autoencoder that takes parameters for perceptual loss (VGG), mse loss, and mae loss, so I can find the perfect weighting for the 3 losses.

When I evaluate my trained model, all parameters produce the same result for a given image reconstruction. The network learned to try to find the best solution for all 3 losses and ignores the parameters it sees.

One important note is that if I train the network with fixed loss parameters (still using YOTO and FiLM, but no random param sampling) I get the expected results for that set of parameters. Everything works fine until I let it take a random sample of params during training.

To condition the network on the params, the paper uses FiLM. You use a fully connected layer that takes the params as input and outputs a scale and shift, then an affine transformation is done on the layer's activations using these two outputs. The output at each layer becomes a * γ + β, where a is the activations of the layer and γ,β are the scale and shift we want to apply (channel wise).

My problem might be in how I do channel wise multiplication between a tensor (the activations) and a matrix (the shift or scale values for each channel in the activations). For example I need to multiply activations with shape (batchSize, 8, 8, 512) by a scaling factor of shape (batchSize, 512). The way I do it is to reshape the matrix to be (batchSize, 1, 1, 512) which allows for a valid multiplication but I'm not sure if the broadcasting is doing what I think it is, or if it's another problem altogether.

Here is my decoder code. For each layer I do transposed conv, batch norm (non trainable as paper specifies), shift and scale to condition on the parameters, then relu.

def FCNet(params, channelCount):
    x = Dense(512)(params)
    x = Activation('relu')(x)
    shift = Dense(channelCount)(x)
    scale = Dense(channelCount)(x)
    return shift, scale

def buildYOTODecoder():
    latent = Input(shape=(8, 8, 512))
    params = Input(shape=(paramCount,))

    depth = [512, 256, 128, 64, 32]
    for d in range(len(depth)):
        if d == 0:
            x = Conv2DTranspose(depth[d], (3, 3), strides = (2, 2), padding = 'same', kernel_initializer=RandomNormal(stddev=0.02))(latent)
        else:
            x = Conv2DTranspose(depth[d], (3, 3), strides = (2, 2), padding = 'same', kernel_initializer=RandomNormal(stddev=0.02))(x)
        x = BatchNormalization(center=False, scale=False)(x)
        m, v = FCNet(params, depth[d])
        rs = [tf.keras.backend.shape(latent)[0], 1, 1, depth[d]]
        m = tf.reshape(m, rs)
        v = tf.reshape(v, rs)
        x = Activation('relu')(x * v + m)

    x = convBlock(3, (3, 3), (1, 1), x, 'none', 'sigmoid')

    model = Model([latent, params], x)
    return model

Here is my loss function where I'm also not sure if the way I'm reshaping, multiplying, and indexing is doing what it should do. I use 7 parameters to weight the loss, one for each VGG layer, then one for MSE and MAE loss

def ParamVGGLoss(params):
    def loss(y_true, y_pred):
        paramsReshaped = tf.reshape(params, [tf.keras.backend.shape(y_true)[0], 1, 1, paramCount])
        true = vgg(preprocess_input(y_true * 255))
        pred = vgg(preprocess_input(y_pred * 255))

        vggLoss = 0
        for i in range(len(true)):
            t = normalize_tensor(true[i])
            p = normalize_tensor(pred[i])
            sqDif = tf.math.square(t - p) * paramsReshaped[:, :, :, i : i + 1]
            vggLoss += tf.math.reduce_mean(sqDif)

        maeLoss = tf.math.abs(y_true - y_pred) * paramsReshaped[:, :, :, 5 : 6]
        maeLoss = tf.math.reduce_mean(maeLoss)
        mseLoss = tf.math.square(y_true - y_pred) * paramsReshaped[:, :, :, 6:]
        mseLoss = tf.math.reduce_mean(mseLoss) 
        return vggLoss + maeLoss + mseLoss
    return loss