Added TI decay start offset

1 files changed, 8 insertions, 2 deletions

diff --git a/train_ti.py b/train_ti.py
index 3b7e3b1..e18ee38 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -415,11 +415,17 @@ def parse_args():
     )
     parser.add_argument(
         "--decay_factor",
-        default=100,
+        default=1,
         type=float,
         help="Embedding decay factor."
     )
     parser.add_argument(
+        "--decay_start",
+        default=1e-4,
+        type=float,
+        help="Embedding decay start offset."
+    )
+    parser.add_argument(
         "--noise_timesteps",
         type=int,
         default=1000,
@@ -833,7 +839,7 @@ def main():
     def on_after_optimize(lr: float):
         text_encoder.text_model.embeddings.normalize(
             args.decay_target,
-            min(1.0, args.decay_factor * lr)
+            min(1.0, max(0.0, args.decay_factor * ((lr - args.decay_start) / (args.learning_rate - args.decay_start))))
         )
 
     loop = partial(