Annealing variable in loss function

[Henning Lange]

I am using the tensorflow backend and I am looking for an elegant way of using an annealing variable in my loss function.

What I am trying to do is balance to losses: a*loss1 + (1-a)*loss2

During training I want a to gradually go from 1 to 0.

My hacky solution at the moment is to recompile the network after each epoch which obviously is slow.

[Henning Lange]

Another possible solution would be to have access to the number of iterations within the loss function.

Does anyone have an idea?

[olafv...@gmail.com]

How about some dirty hack accessing the epoch number as a global variable. Not nice, but poossibly better than what you have know?

Pretty new to Keras and python, so forgive if this is totally beside the point. :-)

[olafv...@gmail.com]

After removing 3 annoying typos what I wanted to say was:

How about some dirty hack accessing the epoch number as a global variable. Not nice, but possibly better than what you have now?

Pretty new to Keras and python, so forgive me if this is totally beside the point. :-)

[Henning Lange]

Well yes, that's what I am trying to achieve. I simply don't know how to achieve that though.

[Henning Lange]

So basically what I have now is this after each epoch:

K.get_session().run(net.T.assign(net.T + 1))

But calling this takes approximately 7s after each iteration. There must be a smarter way of doing this.

[sebastien...@gmail.com]

Did you manage to do that ? I have the same problem when trying to implement KL annealing in VAE objective

[christia...@upr.edu]

All,

Not sure if this is still relevant, but I thought it might serve to anyone trying to implement KL annealing. Here is my solution,

rate = K.variable(0.0,name='KL_Annealing')

annealing_rate = 0.0001

def vae_loss(y_true,y_pred):

    global annealing_rate

    global rate

    xent_loss = keras.losses.categorical_crossentropy(y_true, y_pred)

    kl_loss = -rate * K.mean(1 + output_log - K.square(output_mean) - K.exp(output_log), axis=-1)

    rate = K.tf.assign_add(rate,annealing_rate)

    annealing_rate *= 2

    rate = K.tf.assign(rate,K.clip(rate,0.0,1.0))

    return xent_loss + kl_loss

It's still a bit hacky, but runs smoothly (at least much less than 7 seconds); in fact, my total training time stayed the same. The overall function follows Francois' blog post on AEs.

I hope it helps.

Thanks,

Christian L.

[Christian J Lagares Nieves]

PS I just noticed the rate is extremely high. I updated the function as,

rate = K.variable(0.0,name='KL_Annealing')

annealing_rate = 0.0001

def vae_loss(y_true,y_pred):

   global annealing_rate

   global rate

   xent_loss = keras.losses.categorical_crossentropy(y_true, y_pred)

   kl_loss = -rate * K.mean(1 + output_log - K.square(output_mean) - K.exp(output_log), axis=-1)

   rate = K.tf.assign(rate,annealing_rate)

   annealing_rate *=1.05

   rate = K.tf.assign(rate,K.clip(rate,0.0,1.0))

   return xent_loss + kl_loss