Warmup小記

魚與魚發表於2022-04-10

什麼是warmup

熱身，在剛剛開始訓練時以很小的學習率進行訓練，使得網路熟悉資料，隨著訓練的進行學習率慢慢變大，到了一定程度，以設定的初始學習率進行訓練，接著過了一些inter後，學習率再慢慢變小；

學習率變化：上升——平穩——下降

為什麼用warmup

有助於減緩模型在初始階段對mini-batch的提前過擬合現象，保持分佈的平穩
有助於保持模型深層的穩定性

可以認為，剛開始模型對資料的“分佈”理解為零，或者是說“均勻分佈”（當然這取決於你的初始化）；在第一輪訓練的時候，每個資料點對模型來說都是新的，模型會很快地進行資料分佈修正，如果這時候學習率就很大，極有可能導致開始的時候就對該資料“過擬合”，後面要通過多輪訓練才能拉回來，浪費時間。當訓練了一段時間（比如兩輪、三輪）後，模型已經對每個資料點看過幾遍了，或者說對當前的batch而言有了一些正確的先驗，較大的學習率就不那麼容易會使模型學偏，所以可以適當調大學習率。這個過程就可以看做是warmup。

當模型訓到一定階段後（比如十個epoch），模型的分佈就已經比較固定了，或者說能學到的新東西就比較少了。如果還沿用較大的學習率，就會破壞這種穩定性，用我們通常的話說，就是已經接近loss的local optimal了，為了靠近這個point，我們就要慢慢來。

這裡只摘錄了一小段，參考文獻 [1] 解釋的很好。

learning rate schedule

warmup和learning schedule是類似的，只是學習率變化不同。如圖

learning rate schedule

tensorflow 中有幾種不同的learning rate schedule，以上圖的3種為例，更多schedule可以直達官網

# CosineDecay
cosine_learning_rate_schedule = tf.keras.optimizers.schedules.CosineDecay(0.001,4000)
plt.plot(cosine_learning_rate_schedule(tf.range(40000, dtype=tf.float32)),label="cosine")

# ExponentialDecay
exp_learning_rate_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
    0.001, 4000, 0.9, staircase=False, name=None
)
plt.plot(exp_learning_rate_schedule(tf.range(40000, dtype=tf.float32)),label="exp")

# PiecewiseConstantDecay
boundaries = [10000, 20000,30000]
values = [0.001, 0.0008, 0.0004,0.0001]
piecewise_learning_rate_schedule = tf.keras.optimizers.schedules.PiecewiseConstantDecay(
    boundaries, values)
plt.plot([piecewise_learning_rate_schedule(step) for step in tf.range(40000, dtype=tf.float32)],label="piecewise")

# 自定義 Schedule
my_learning_rate_schedule = MySchedule(0.001)
plt.plot([my_learning_rate_schedule(step) for step in tf.range(40000, dtype=tf.float32)],label="warmup")

plt.title("Learning rate schedule")
plt.ylabel("Learning Rate")
plt.xlabel("Train Step")
plt.legend()

# 自定義 Schedule
class MySchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
  def __init__(self, initial_learning_rate, warmup_steps=4000):
    super(MySchedule, self).__init__()
    self.initial_learning_rate = initial_learning_rate
    self.warmup_steps = warmup_steps
    
  def __call__(self, step):
    if step > self.warmup_steps:
      return self.initial_learning_rate * self.warmup_steps * step ** -1
    else:
      return self.initial_learning_rate * step * (self.warmup_steps ** -1)

warmup in transformer

Noam Optimizer

\[\alpha \frac{1}{\sqrt{d_{model}}}min(\frac{1}{\sqrt{t}},\frac{t}{w^{3/2}}) \]

\[lrate = d^{-0.5}_{model}*min(step\_ num^{-0.5},step\_ num*warmup\_ steps^{-1.5}) \]

class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
  def __init__(self, d_model, warmup_steps=4000):
    super(CustomSchedule, self).__init__()
    
    self.d_model = d_model
    self.d_model = tf.cast(self.d_model, tf.float32)

    self.warmup_steps = warmup_steps
    
  def __call__(self, step):
    arg1 = tf.math.rsqrt(step)
    arg2 = step * (self.warmup_steps ** -1.5)
    
    return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)

learning_rate = CustomSchedule(d_model)
optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98, 
                                     epsilon=1e-9)

temp_learning_rate_schedule = CustomSchedule(128)

plt.plot(temp_learning_rate_schedule(tf.range(40000, dtype=tf.float32)))
plt.ylabel("Learning Rate")
plt.xlabel("Train Step")