250 lines
10 KiB
Lua
250 lines
10 KiB
Lua
-- ============================================================================
|
|
-- SPDX-License-Identifier: GPL-3.0-or-later
|
|
-- Copyright (C) 2026 Alexander Allan (MDMAchine) -- A&E Concepts
|
|
-- ============================================================================
|
|
|
|
-- MD Dual-Time V2 -- Inner Convergence Loop Sampler
|
|
-- MDMAchine | A&E Concepts (c) 2026
|
|
--
|
|
-- Zero-NFE iterative refinement using cached (v, sigma, x) tuples.
|
|
-- Inverse-distance + sigma-proximity velocity interpolation. Sigma-adaptive
|
|
-- inner blend protects early vocal separation. owns_loop = true. Single NFE.
|
|
-- ============================================================================
|
|
|
|
local C = require("md_solver_commons")
|
|
|
|
-- ── VELOCITY INTERPOLATION (per-batch) ──────────────────────────────────────
|
|
|
|
local function interpolate_velocity_batch(x_cand, off, cnt, cache, cache_len, sigma_curr, sigma_wt)
|
|
local v_interp = {}
|
|
for i = 0, cnt - 1 do v_interp[i] = 0.0 end
|
|
local total_weight = 0.0
|
|
|
|
for k = 1, cache_len do
|
|
local entry = cache[k]
|
|
local dist_sq = 0.0
|
|
for i = 0, cnt - 1 do
|
|
local d = x_cand[off + i] - entry.x[off + i]
|
|
dist_sq = dist_sq + d * d
|
|
end
|
|
local pos_dist = math.sqrt(dist_sq / math.max(cnt, 1) + C.EPSILON)
|
|
local sig_factor = 1.0 / (1.0 + sigma_wt * math.abs(sigma_curr - entry.sigma))
|
|
local w = sig_factor / (pos_dist + C.EPSILON)
|
|
total_weight = total_weight + w
|
|
for i = 0, cnt - 1 do v_interp[i] = v_interp[i] + w * entry.v[off + i] end
|
|
end
|
|
|
|
if total_weight > C.EPSILON then
|
|
local inv_w = 1.0 / total_weight
|
|
for i = 0, cnt - 1 do v_interp[i] = v_interp[i] * inv_w end
|
|
end
|
|
return v_interp
|
|
end
|
|
|
|
-- ── INNER CONVERGENCE LOOP ──────────────────────────────────────────────────
|
|
|
|
local function inner_loop(x_start, x_euler, v_curr, dt, n, B, NPB,
|
|
cache, cache_len, sigma_curr, sigma_wt,
|
|
max_inner, conv_thresh, relaxation)
|
|
local x_cand = C.vec_clone(x_euler, n)
|
|
local init_resid, final_resid = 0.0, 0.0
|
|
local converged = false
|
|
local iters_used = 0
|
|
|
|
for k = 1, max_inner do
|
|
iters_used = k
|
|
|
|
local v_interp_full = C.vec_clone(v_curr, n)
|
|
for b = 0, B - 1 do
|
|
local off = b * NPB
|
|
local v_batch = interpolate_velocity_batch(
|
|
x_cand, off, NPB, cache, cache_len, sigma_curr, sigma_wt)
|
|
for j = 0, NPB - 1 do v_interp_full[off + j] = v_batch[j] end
|
|
end
|
|
|
|
local x_new = {}
|
|
for j = 0, n - 1 do x_new[j] = x_start[j] + dt * v_interp_full[j] end
|
|
|
|
local corr_rms = 0.0
|
|
for j = 0, n - 1 do
|
|
local c = x_new[j] - x_cand[j]
|
|
corr_rms = corr_rms + c * c
|
|
end
|
|
corr_rms = math.sqrt(corr_rms / math.max(n, 1))
|
|
|
|
if k == 1 then init_resid = corr_rms end
|
|
final_resid = corr_rms
|
|
|
|
for j = 0, n - 1 do
|
|
x_cand[j] = x_cand[j] + relaxation * (x_new[j] - x_cand[j])
|
|
end
|
|
|
|
if C.has_nan_inf(x_cand, n) then
|
|
for j = 0, n - 1 do x_cand[j] = x_euler[j] end
|
|
break
|
|
end
|
|
|
|
if corr_rms < conv_thresh then converged = true; break end
|
|
end
|
|
|
|
return x_cand, iters_used, converged, init_resid, final_resid
|
|
end
|
|
|
|
-- ── SOLVER DEFINITION ───────────────────────────────────────────────────────
|
|
|
|
solver = {
|
|
name = "md_dual_time_v2",
|
|
display = "MD Dual-Time V2",
|
|
description = "Inner convergence loop sampler. Zero-NFE velocity history interpolation. Sigma-adaptive inner blend. Batch-aware, shared anchor stack.",
|
|
nfe = 1,
|
|
order = 1,
|
|
needs_model = false,
|
|
stateful = true,
|
|
stochastic = true,
|
|
owns_loop = true,
|
|
params = {
|
|
{ key = "max_inner", type = "slider", label = "Max Inner Iterations",
|
|
default = 3, min = 1, max = 10, step = 1,
|
|
hint = "Pseudo-time iterations per step." },
|
|
{ key = "cache_depth", type = "slider", label = "Cache Depth",
|
|
default = 6, min = 2, max = 12, step = 1,
|
|
hint = "Number of (v, sigma, x) tuples stored." },
|
|
{ key = "convergence_threshold", type = "slider", label = "Convergence Threshold",
|
|
default = 0.005, min = 0.0005, max = 0.1, step = 0.0005,
|
|
hint = "Per-element RMS for early exit." },
|
|
{ key = "sigma_weight", type = "slider", label = "Sigma Proximity Weight",
|
|
default = 4.0, min = 0.0, max = 8.0, step = 0.25,
|
|
hint = "Sigma proximity influence. Higher = less phase ghosting." },
|
|
{ key = "relaxation", type = "slider", label = "Relaxation Factor",
|
|
default = 0.45, min = 0.1, max = 1.0, step = 0.05,
|
|
hint = "Inner loop step size. Lower = less phase ghosting." },
|
|
{ key = "sigma_gate", type = "slider", label = "Sigma Gate",
|
|
default = 0.9, min = 0.5, max = 1.0, step = 0.05,
|
|
hint = "Sigma fraction above which inner loop is disabled." },
|
|
{ key = "inner_blend", type = "slider", label = "Inner Blend",
|
|
default = 0.4, min = 0.0, max = 1.0, step = 0.05,
|
|
hint = "Max Euler/converged mix. Sigma-adaptive: near-zero early, ramps quadratically." },
|
|
},
|
|
}
|
|
|
|
C.append_common_params(solver.params)
|
|
|
|
-- ── SAMPLE ──────────────────────────────────────────────────────────────────
|
|
|
|
function sample(xt, vt_buf, schedule, n, model_fn)
|
|
local p = params or {}
|
|
local B, NPB = C.get_batch_routing(n)
|
|
|
|
local max_inner = math.floor(C.num_param(p, "max_inner", 3))
|
|
local cache_depth = math.floor(C.num_param(p, "cache_depth", 6))
|
|
local conv_thresh = C.num_param(p, "convergence_threshold", 0.005)
|
|
local sigma_wt = C.num_param(p, "sigma_weight", 4.0)
|
|
local relaxation = C.num_param(p, "relaxation", 0.45)
|
|
local sigma_gate = C.num_param(p, "sigma_gate", 0.9)
|
|
local inner_blend = C.num_param(p, "inner_blend", 0.4)
|
|
local opts = C.read_common_opts(p)
|
|
local state = C.new_state()
|
|
|
|
-- Engine schedule has NO trailing 0 (fix ported from 46c081e): iterate all ns
|
|
-- entries so the last iteration gets sigma_next = 0.0 and the terminal branch
|
|
-- performs the final x0 projection. With ns - 1 that branch is dead code and
|
|
-- the output keeps ~final-sigma noise.
|
|
local ns, n_steps = #schedule, #schedule
|
|
if n_steps < 1 then return end
|
|
|
|
local sigma_max = schedule[1]
|
|
local cache, cache_len, cache_pos, cache_max = {}, 0, 0, cache_depth
|
|
for k = 1, cache_depth do cache[k] = nil end
|
|
|
|
local x = C.fa_to_tbl(xt, n)
|
|
|
|
if opts.verbose then
|
|
print(string.format("[DUAL-TIME V2] Schedule: %d steps | B=%d NPB=%d | inner=%d cache=%d",
|
|
n_steps, B, NPB, max_inner, cache_depth))
|
|
end
|
|
|
|
for i = 1, n_steps do
|
|
local sigma_curr = schedule[i]
|
|
local sigma_next = (i < ns) and schedule[i + 1] or 0.0
|
|
local step_idx = i - 1
|
|
local sigma_ratio = C.clamp(sigma_curr / math.max(sigma_max, C.EPSILON), 0.0, 1.0)
|
|
|
|
if sigma_next == 0.0 then
|
|
C.tbl_to_fa(x, xt, n)
|
|
model_fn(xt, sigma_curr)
|
|
local v_final = C.fa_to_tbl(vt_buf, n)
|
|
for j = 0, n - 1 do x[j] = x[j] - v_final[j] * sigma_curr end
|
|
break
|
|
end
|
|
|
|
C.tbl_to_fa(x, xt, n)
|
|
model_fn(xt, sigma_curr)
|
|
local v_curr = C.fa_to_tbl(vt_buf, n)
|
|
local dt = sigma_next - sigma_curr
|
|
|
|
-- Relational decomposition
|
|
if opts.rw > 0 then
|
|
C.apply_relational(v_curr, n, B, NPB, sigma_ratio, sigma_max,
|
|
opts.rw, opts.rw_sigma_pow, opts.drift_on, opts.drift_thr, x)
|
|
end
|
|
|
|
local x_euler = {}
|
|
for j = 0, n - 1 do x_euler[j] = x[j] + dt * v_curr[j] end
|
|
|
|
local x_new = x_euler
|
|
local iters_used, converged, init_resid, final_resid = 0, false, 0.0, 0.0
|
|
|
|
if cache_len >= 2 and sigma_ratio < sigma_gate then
|
|
local cache_ordered = {}
|
|
for k = 1, cache_len do
|
|
local idx = ((cache_pos - cache_len + k - 1) % cache_max) + 1
|
|
cache_ordered[k] = cache[idx]
|
|
end
|
|
|
|
x_new, iters_used, converged, init_resid, final_resid = inner_loop(
|
|
x, x_euler, v_curr, dt, n, B, NPB,
|
|
cache_ordered, cache_len, sigma_curr, sigma_wt,
|
|
max_inner, conv_thresh, relaxation)
|
|
|
|
-- Sigma-adaptive inner blend
|
|
local blend_ramp = (1.0 - sigma_ratio) * (1.0 - sigma_ratio)
|
|
local eff_blend = inner_blend * blend_ramp
|
|
if eff_blend > 1e-6 and eff_blend < 1.0 - 1e-6 then
|
|
for j = 0, n - 1 do
|
|
x_new[j] = (1.0 - eff_blend) * x_euler[j] + eff_blend * x_new[j]
|
|
end
|
|
elseif eff_blend <= 1e-6 then
|
|
for j = 0, n - 1 do x_new[j] = x_euler[j] end
|
|
end
|
|
end
|
|
|
|
-- Cache push
|
|
cache_pos = (cache_pos % cache_max) + 1
|
|
cache[cache_pos] = { v = C.vec_clone(v_curr, n), sigma = sigma_curr, x = C.vec_clone(x, n) }
|
|
if cache_len < cache_max then cache_len = cache_len + 1 end
|
|
|
|
if C.has_nan_inf(x_new, n) then
|
|
for j = 0, n - 1 do x_new[j] = x_euler[j] end
|
|
end
|
|
|
|
opts.sigma_next = sigma_next
|
|
opts.step_idx = step_idx
|
|
C.post_advance(x_new, n, B, NPB, sigma_ratio, opts, state)
|
|
|
|
if opts.verbose then
|
|
print(string.format("[DUAL-TIME V2] step %02d | inner=%d/%d %s | resid %.5f->%.5f | rms=%.3f",
|
|
step_idx, iters_used, max_inner,
|
|
converged and "CONV" or (iters_used > 0 and "max" or "skip"),
|
|
init_resid, final_resid, C.rms(x_new, n)))
|
|
end
|
|
|
|
x = x_new
|
|
C.tbl_to_fa(x, xt, n)
|
|
C.tbl_to_fa(v_curr, vt_buf, n)
|
|
if on_step(step_idx, sigma_curr, sigma_next) then return end
|
|
x = C.fa_to_tbl(xt, n)
|
|
end
|
|
|
|
C.tbl_to_fa(x, xt, n)
|
|
end
|