From cfdce99957bf99d6fec61bd939c4cbc3e02b3b83 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Fri, 10 Apr 2026 00:08:18 +0100 Subject: [PATCH] fix ER-SDE flow-matching parameterization using alpha=1 The VP update formula uses lambda=sigma/alpha which blows up to ~1e8 for flow matching (where alpha=1-sigma approaches 0). This makes r_fn=fn(next)/fn(curr) collapse to ~0, causing the sampler to discard the previous sample and independently re-predict x0 at each step. Set alpha=1 and lambda=sigma for flow matching so the VP formula naturally reduces to the correct form r_f*x + (1-r_f)*x0 with lambda in [0,1] where the noise function produces meaningful ratios. --- modules/schedulers/scheduler_ersde.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/modules/schedulers/scheduler_ersde.py b/modules/schedulers/scheduler_ersde.py index df58a78ee..a81fe14df 100644 --- a/modules/schedulers/scheduler_ersde.py +++ b/modules/schedulers/scheduler_ersde.py @@ -201,7 +201,12 @@ class ERSDEScheduler(SchedulerMixin, ConfigMixin): self.lower_order_nums = 0 def _setup_flow(self, sigmas, device, mu=None): - """Common flow-matching schedule setup from a sigmas tensor.""" + """Common flow-matching schedule setup from a sigmas tensor. + + Uses alpha=1 so the VP update formula r_alpha*r_fn*x + alpha*(1-r_fn)*x0 + reduces to the correct flow-matching form r_f*x + (1-r_f)*x0, and lambda=sigma + keeps the customized noise function in a well-behaved range on [0,1]. + """ self._is_flow = True if self.config.use_dynamic_shifting: if mu is None: @@ -210,8 +215,8 @@ class ERSDEScheduler(SchedulerMixin, ConfigMixin): else: sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas) flow_sigmas = sigmas.clamp(min=1e-8, max=1.0 - 1e-8) - flow_alphas = 1.0 - flow_sigmas - flow_lambdas = flow_sigmas / flow_alphas + flow_alphas = torch.ones_like(flow_sigmas) + flow_lambdas = flow_sigmas.clone() self.timesteps = (sigmas * self.config.num_train_timesteps).to(device=device) self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device, dtype=sigmas.dtype)]) self._flow_alphas = torch.cat([flow_alphas, torch.ones(1, device=device, dtype=torch.float64)])