import math


def decay_annealing(train_step: int, final_train_step: int,
                    beta_final: float = 1.0, beta_start: float = 0.0) -> float:
    """Exponential decay annealing such that at final_train_step and from there on beta is beta_final."""
    if train_step < final_train_step:
        decay_rate = math.log(beta_final / beta_start) / final_train_step
        return beta_start * math.exp(decay_rate * train_step)
    else:
        return beta_final