# Custom Learning Rate

**URL:** <https://discuss.bayesflow.org/t/custom-learning-rate/64>\
**Category:** General\
**Created:** [February 19, 2024, 2:51am UTC](https://discuss.bayesflow.org/t/custom-learning-rate/64 "2024-02-19T02:51:41Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![ali](https://avatars.discourse-cdn.com/v4/letter/a/8c91f0/32.png) [@ali](https://discuss.bayesflow.org/u/ali)\
**Post date:** [February 19, 2024, 2:51am UTC](https://discuss.bayesflow.org/t/custom-learning-rate/64/1 "2024-02-19T02:51:41Z")

</div>

Hi,

We have found that resetting the learning rate and introducing new data (and discarding the old one) in discrete rounds yields better performance for our model. Is it possible to have a similar mechanism in round-based training to have a learning rate with cosine decay in each round? Or, is it possible to define a custom learning rate for the manual discrete rounds method so that in each round, the cosine decay starts from a lower initial learning rate (because the model has already learned some patterns from the earlier rounds)?

Thank you very much,  
Ali

PS. Having the both options is extremely helpful. Also, I tried to experiment with the following but got an error suggesting a conflict with the cosine decay:

lr\_schedule = tf.keras.optimizers.schedules.ExponentialDecay(  
initial\_learning\_rate,  
decay\_steps=decay\_steps,  
decay\_rate=decay\_rate,  
staircase=True)

---

<div class="post-metadata">

**Author:** ![elseml](https://yyz1.discourse-cdn.com/flex007/user_avatar/discuss.bayesflow.org/elseml/32/19_2.png) [@elseml](https://discuss.bayesflow.org/u/elseml)\
**Post date:** [February 19, 2024, 1:58pm UTC](https://discuss.bayesflow.org/t/custom-learning-rate/64/2 "2024-02-19T13:58:57Z")

</div>

Hi Ali,

if I understand you correctly, a [cosine decay schedule with restarts](https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/schedules/CosineDecayRestarts) could be helpful for your situation. You can implement it in your training like this:

```auto
schedule = tf.keras.optimizers.schedules.CosineDecayRestarts(
    initial_learning_rate=0.0005,
    first_decay_steps=1000,
    t_mul=1.0,
    m_mul=0.5,
    alpha=0.0
)
optimizer = tf.keras.optimizers.Adam(schedule, global_clipnorm=1.0)
history = trainer.train_rounds([YOUR CODE], optimizer=optimizer)

```

Edit: I just saw that you are discarding the old data, so t\_mul can stay at 1.0 and you can adapt m\_mul to control the decrease of the learning rate over restarts.  
If you want use train\_rounds and keep previously simulated data, the growing number of data sets makes this a bit tricky, since the multiplicate factor t\_mul cannot match the additive growth of the number of iterations within the rounds. You could allow for such a behavior by constructing a [custom schedule](https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/schedules/LearningRateSchedule), possibly building upon the [TensorFlow CosineDecayRestarts code](https://github.com/keras-team/keras/blob/v2.15.0/keras/optimizers/schedules/learning_rate_schedule.py#L770-L911).

Cheers,  
Lasse

---

<div class="post-metadata">

**Author:** ![KLDivergence](https://yyz1.discourse-cdn.com/flex007/user_avatar/discuss.bayesflow.org/kldivergence/32/15_2.png) [@KLDivergence](https://discuss.bayesflow.org/u/KLDivergence)\
**Post date:** [February 19, 2024, 3:05pm UTC](https://discuss.bayesflow.org/t/custom-learning-rate/64/3 "2024-02-19T15:05:54Z")

</div>

We could add a flag for discarding / keeping old data from previous rounds in round-based training? I assume you are currently “hacking” it with an external loop over offline training?

---

<div class="post-metadata">

**Author:** ![ali](https://avatars.discourse-cdn.com/v4/letter/a/8c91f0/32.png) [@ali](https://discuss.bayesflow.org/u/ali)\
**Post date:** [February 19, 2024, 3:28pm UTC](https://discuss.bayesflow.org/t/custom-learning-rate/64/4 "2024-02-19T15:28:58Z")

</div>

Thank you very much, both. Yes, I am “hacking” it with an external loop over offline training. The only reason is to avoid high computational costs in the last rounds while providing large batches of data in each round. I think there’s a tradeoff between the two methods. I feel this approach helps to adaptively focus the model’s learning on areas of the parameter space that are poorly understood, whereas the current round-based training refines the posterior approximation with increasing precision, which might not be achieved otherwise.