adding `beta_1` parameter to `create_optimizer` in `optimization.py` to allow...

adding `beta_1` parameter to `create_optimizer` in `optimization.py` to allow for a tunable `beta_1`. Defaults to the original `0.9` to ensure existing compatibility.

PiperOrigin-RevId: 365634937

official/nlp/optimization.py

@@ -69,7 +69,8 @@ def create_optimizer(init_lr,
                      num_train_steps,
                      num_warmup_steps,
                      end_lr=0.0,
-                     optimizer_type='adamw'):
+                     optimizer_type='adamw',
+                     beta_1=0.9):
   """Creates an optimizer with learning rate schedule."""
   # Implements linear decay of the learning rate.
   lr_schedule = tf.keras.optimizers.schedules.PolynomialDecay(
@@ -87,7 +88,7 @@ def create_optimizer(init_lr,
     optimizer = AdamWeightDecay(
         learning_rate=lr_schedule,
         weight_decay_rate=0.01,
-        beta_1=0.9,
+        beta_1=beta_1,
         beta_2=0.999,
         epsilon=1e-6,
         exclude_from_weight_decay=['LayerNorm', 'layer_norm', 'bias'])
@@ -96,7 +97,7 @@ def create_optimizer(init_lr,
     optimizer = tfa_optimizers.LAMB(
         learning_rate=lr_schedule,
         weight_decay_rate=0.01,
-        beta_1=0.9,
+        beta_1=beta_1,
         beta_2=0.999,
         epsilon=1e-6,
         exclude_from_weight_decay=['LayerNorm', 'layer_norm', 'bias'])