examples/research_projects/bertabs/modeling_bertabs.py/BertSumOptimizer/_update_rate
class BertSumOptimizer:    def _update_rate(self, stack):
        return self.lr[stack] * min(self._step ** (-0.5), self._step * self.warmup_steps[stack] ** (-1.5))