Is there a way to pass along temporal weights to a loss function?

Question

Background

Currently, I'm using an LSTM to perform a regression. I'm using small batch sizes with a reasonably large amount of timesteps (but much, much fewer than the number of timesteps I have).

I'm attempting to transition to larger batches with fewer timesteps, but with stateful enabled, to allow a larger amount of generated training data to be used.

However, I am currently using a regularization based off of sqrt(timestep), (this is ablation tested and helps with convergence speed, it works because of the statistical nature of the problem, expected error decreases by a factor of sqrt(timestep)). This is performed by using tf.range to generate a list of the proper size within the loss function. This approach will not be correct when stateful is enabled, because it will be counting the wrong number of timesteps (number of timesteps in this batch, rather than seen so far overall).

Question

Is there a way to pass an offset or list of ints or floats to the loss function? Preferably without modifying the model, but I recognize that a hack of this nature might be required.

Code

Simplified model:

def create_model():    
    inputs = Input(shape=(None,input_nodes))
    next_input = inputs
    for i in range(dense_layers):
        dense = TimeDistributed(Dense(units=dense_nodes,
                activation='relu',
                kernel_regularizer=l2(regularization_weight),
                activity_regularizer=l2(regularization_weight)))\
            (next_input)
        next_input = TimeDistributed(Dropout(dropout_dense))(dense)

    for i in range(lstm_layers):
        prev_input = next_input
        next_input = LSTM(units=lstm_nodes,
                dropout=dropout_lstm,
                recurrent_dropout=dropout_lstm,
                kernel_regularizer=l2(regularization_weight),
                recurrent_regularizer=l2(regularization_weight),
                activity_regularizer=l2(regularization_weight),
                stateful=True,
                return_sequences=True)\
            (prev_input)
        next_input = add([prev_input, next_input])

    outputs = TimeDistributed(Dense(output_nodes,
            kernel_regularizer=l2(regularization_weight),
            activity_regularizer=l2(regularization_weight)))\
        (next_input)

    model = Model(inputs=inputs, outputs=outputs)

Loss function

def loss_function(y_true, y_pred):
    length = K.shape(y_pred)[1]

    seq = K.ones(shape=(length,))
    if use_sqrt_loss_scaling:
        seq = tf.range(1, length+1, dtype='int32')
        seq = K.sqrt(tf.cast(seq, tf.float32))

    seq = K.reshape(seq, (-1, 1))

    if separate_theta_phi:
        angle_loss = phi_loss_weight * phi_metric(y_true, y_pred, angle_loss_fun)
        angle_loss += theta_loss_weight * theta_metric(y_true, y_pred, angle_loss_fun)
    else:
        angle_loss = angle_loss_weight * total_angle_metric(y_true, y_pred, angle_loss_fun)

    norm_loss = norm_loss_weight * norm_loss_fun(y_true, y_pred)
    energy_loss = energy_loss_weight * energy_metric(y_true, y_pred)
    stability_loss = stability_loss_weight * stab_loss_fun(y_true, y_pred)
    act_loss = act_loss_weight * act_loss_fun(y_true, y_pred)

    return K.sum(K.dot(0
        + angle_loss
        + norm_loss
        + energy_loss
        + stability_loss
        + act_loss
        , seq))

(The functions that calculate the pieces of the loss function shouldn't be super duper relevant. Simply, they're also loss functions.)

I am not sure I understand your aim correctly. But is this offset you want to pass to the loss function batch-dependent? i.e. is there a different offset per batch that you can pass it alongside to the model as input? — today, Apr 15 '20 at 01:03
Exactly! Like, if I'm looking at timesteps (0, 100) I'd like to be able to calculate the square roots of each integer from 0 to 100 and dot it with my raw loss values. But if I'm looking at the 4th batch of timesteps, I'll be looking at timesteps (300, 400), so I want to take the sqrt of the numbers from 300 to 400 and dot it with my raw loss values. — OmnipotentEntity, Apr 15 '20 at 01:06

today · Accepted Answer · 2020-04-15T01:28:43.727

For that purpose, you can use sample_weight argument of fit method and pass sample_weight_mode='temporal' to compile method so that you can assign a weight to each timestep of each sample in the batch:

model.compile(..., sample_weight_mode='temporal')
model.fit(..., sample_weight=sample_weight)

The sample_weight should be an array of shape (num_samples, num_timesteps).

Note that if you are using an input data generator or an instance of Sequence, instead you need to pass the sample weights as the third element of generated tuple/list in the generator or Sequence instance.