End of training

lr_scheduler/lr_scheduler.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:cc9a5d647349c541df37221f857f1528d69a64ee8921808d7169158f833bb344
+oid sha256:b7ff260f2e05fbce740edbdb45e6fec95b731941ad4077b051c8da5bb1f684a1
 size 1076

optimizer/optimizer.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:50f22f9265e94e528478207e4fc002476ad0034abeea03c61fa14773b8587fff
+oid sha256:60d2ec204528438aa57b245329bdb08da454a36321dac7fa203011a06d282a60
 size 1149290878

scheduler/scheduler_config.py

@@ -23,21 +23,23 @@ class SDEPolynomialConfig:
     drift_degree = 20
     diffusion_degree = 20
 
-    drift_parameters = Matrix([sympy.symbols(f"f:{drift_degree}", real=True, nonzero=True)])
+    #drift_parameters = Matrix([sympy.symbols(f"f:{drift_degree}", real=True, nonzero=True)])
+    drift_parameters = Matrix([sympy.symbols(f"f0", real=True)])
 
-    diffusion_parameters = Matrix([sympy.symbols("l0", real=True, nonzero=True)])
+    diffusion_parameters = Matrix([sympy.symbols(f"l:{diffusion_degree}", real=True, nonzero=True)])
 
 
     @property
     def drift(self): 
         transformed_variable = self.variable
-        return -sympy.Abs(sum(sympy.HadamardProduct(Matrix([[transformed_variable**i for i in range(1,self.drift_degree+1)]]), self.drift_parameters).doit()))
+        return - (transformed_variable**2 + 4 * transformed_variable**4 + 2* transformed_variable**5)
+        #return -sympy.Abs(sum(sympy.HadamardProduct(Matrix([[transformed_variable**i for i in range(1,self.drift_degree+1)]]), self.drift_parameters).doit()))
         
 
     @property
     def diffusion(self):
-        
-        return self.variable**(self.diffusion_parameters[0]**2)
+        #return self.variable * 1e-10
+        return self.variable**(sum(sympy.HadamardProduct(Matrix([[self.variable**i for i in range(0,self.diffusion_degree)]]),self.diffusion_parameters.applyfunc(lambda x: x**2)).doit()))
 
     # TODO (KLAUS) : in the SDE SAMPLING CHANGING Q impacts how we sample z ~ N(0, Q*(delta t))
     diffusion_matrix = 1 
@@ -52,4 +54,4 @@ class SDEPolynomialConfig:
 
 
     target = "epsilon" # x0
-    non_symbolic_parameters = {'drift': torch.ones(drift_degree), 'diffusion': torch.tensor([1.])}
+    non_symbolic_parameters = {'drift': torch.tensor([0.]), 'diffusion': torch.ones(diffusion_degree)}

scheduler/sdeparameters.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:eeb1459bbcbc570e1500ace6bd35966cdd2d7991ab5ad91635bdf3d0080d7acf
+oid sha256:6995386ca1a48ac6d10982761fc3bf36ccb62a39c091a042e2883a6daeca94c3
 size 220

unet/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:639d03b452944d97c7f29345d8d18affc241952a3f70890cc923570fa28e5448
+oid sha256:45111d898ef213981165be7820ec0b934501a727b723f745f824548e03c0a3b8
 size 574476604