1 files changed, 4 insertions, 12 deletions

diff --git a/train_ti.py b/train_ti.py
index 0891c49..fc34d27 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -159,7 +159,7 @@ def parse_args():
     parser.add_argument(
         "--tag_dropout",
         type=float,
-        default=0,
+        default=0.1,
         help="Tag dropout probability.",
     )
     parser.add_argument(
@@ -406,18 +406,12 @@ def parse_args():
         help="Embedding decay target."
     )
     parser.add_argument(
-        "--emb_decay_factor",
-        default=1,
+        "--emb_decay",
+        default=1e-1,
         type=float,
         help="Embedding decay factor."
     )
     parser.add_argument(
-        "--emb_decay_start",
-        default=0,
-        type=float,
-        help="Embedding decay start offset."
-    )
-    parser.add_argument(
         "--noise_timesteps",
         type=int,
         default=1000,
@@ -587,12 +581,10 @@ def main():
         tokenizer=tokenizer,
         sample_scheduler=sample_scheduler,
         checkpoint_output_dir=checkpoint_output_dir,
-        learning_rate=args.learning_rate,
         gradient_checkpointing=args.gradient_checkpointing,
         use_emb_decay=args.use_emb_decay,
         emb_decay_target=args.emb_decay_target,
-        emb_decay_factor=args.emb_decay_factor,
-        emb_decay_start=args.emb_decay_start,
+        emb_decay=args.emb_decay,
         use_ema=args.use_ema,
         ema_inv_gamma=args.ema_inv_gamma,
         ema_power=args.ema_power,