281 lines
11 KiB
Lua
281 lines
11 KiB
Lua
-- ============================================================================
|
|
-- SPDX-License-Identifier: GPL-3.0-or-later
|
|
-- Copyright (C) 2026 Alexander Allan (MDMAchine) -- A&E Concepts
|
|
-- ============================================================================
|
|
|
|
-- MD OmniRelational Solver v3.0 -- Barbour Best Matching + Sigma Decay + Look-Back
|
|
-- MDMAchine | A&E Concepts 2026
|
|
--
|
|
-- WHAT IS NEW IN V3:
|
|
--
|
|
-- V1/V2 applied the same relational weight (rw) at every step regardless
|
|
-- of where you are in the denoising schedule. This is suboptimal:
|
|
--
|
|
-- High sigma (early steps): structure formation phase. The latent is still
|
|
-- mostly noise. Shape-preserving relational geometry matters most here --
|
|
-- normalizing direction prevents any single component dominating.
|
|
--
|
|
-- Low sigma (late steps): detail refinement phase. The latent is close to
|
|
-- x0. Raw velocity is more accurate. Relational re-injection here over-
|
|
-- smooths fine detail.
|
|
--
|
|
-- SIGMA-ADAPTIVE RELATIONAL WEIGHT:
|
|
-- rw_eff = rw * (t_curr / sigma_max) ^ sigma_power
|
|
-- At t=sigma_max (first step): rw_eff = rw (full effect)
|
|
-- At t=0 (final step): rw_eff = 0 (pure Euler)
|
|
-- sigma_power controls the decay curve. 1.0 = linear, 2.0 = quadratic.
|
|
--
|
|
-- LOOK-BACK SNR SMOOTHER (arXiv:2602.09449):
|
|
-- lambda_eff = lb_lambda * (t_curr / sigma_max) ^ lb_snr_power
|
|
-- x_smooth = (1 - lambda_eff) * x_next + lambda_eff * x_prev
|
|
-- Heavy at high sigma (blends structure), fades at low sigma (preserves
|
|
-- detail). Same mechanism as STORM and PingPong. Off by default.
|
|
--
|
|
-- GENERATION STATE RESET:
|
|
-- Module-level state (sigma_max, x_prev for look-back) resets on
|
|
-- step_idx_==0 so same-size consecutive generations don't bleed.
|
|
--
|
|
-- DRIFT GUARD kept from V1 -- simple, clean, no hoisted buffer.
|
|
|
|
solver = {
|
|
name = "md_omni_relational_V3",
|
|
display = "MD OmniRelational V3 (Sigma Adaptive)",
|
|
description = "Barbour Best Matching with sigma-adaptive relational weight and look-back smoother. rw fades toward pure Euler at low sigma. Proper state reset between generations.",
|
|
nfe = 1,
|
|
order = 1,
|
|
needs_model = false,
|
|
stateful = true,
|
|
stochastic = true,
|
|
params = {
|
|
{
|
|
key = "relational_weight",
|
|
type = "slider",
|
|
label = "Relational Weight",
|
|
default = 0.5,
|
|
min = 0.0,
|
|
max = 1.0,
|
|
step = 0.05,
|
|
hint = "Base relational weight at high sigma. Fades toward 0 at low sigma. 0=pure Euler always. 0.5=balanced at structure phase.",
|
|
},
|
|
{
|
|
key = "sigma_power",
|
|
type = "slider",
|
|
label = "Sigma Decay Power",
|
|
default = 1.0,
|
|
min = 0.25,
|
|
max = 4.0,
|
|
step = 0.25,
|
|
hint = "Controls how fast rw fades with sigma. 1.0=linear decay. 2.0=quadratic (faster fade). 0.5=slow fade. Higher = relational effect concentrated earlier.",
|
|
},
|
|
{
|
|
key = "look_back_lambda",
|
|
type = "slider",
|
|
label = "Look-Back Lambda",
|
|
default = 0.0,
|
|
min = 0.0,
|
|
max = 0.5,
|
|
step = 0.01,
|
|
hint = "Look-back coherence smoother. 0=off. Blends current step with previous, fading out at low sigma. Suppresses trajectory shear. Start at 0.05-0.15.",
|
|
},
|
|
{
|
|
key = "look_back_snr_power",
|
|
type = "slider",
|
|
label = "Look-Back SNR Power",
|
|
default = 1.5,
|
|
min = 0.5,
|
|
max = 3.0,
|
|
step = 0.1,
|
|
hint = "Controls how fast look-back fades with sigma. Higher = smoothing concentrated on early structure steps only.",
|
|
},
|
|
{
|
|
key = "drift_guard",
|
|
type = "toggle",
|
|
label = "Drift Guard (AOS)",
|
|
default = false,
|
|
hint = "Project shape_vec onto orthogonal complement of x when cos_sim exceeds threshold. Prevents update reinforcing existing latent structure.",
|
|
},
|
|
{
|
|
key = "drift_threshold",
|
|
type = "slider",
|
|
label = "Drift Threshold",
|
|
default = 0.85,
|
|
min = 0.1,
|
|
max = 1.0,
|
|
step = 0.05,
|
|
hint = "Cosine similarity ceiling before Gram-Schmidt projection fires. 0.85=standard. Only active when Drift Guard is on.",
|
|
},
|
|
{
|
|
key = "eta",
|
|
type = "slider",
|
|
label = "Eta (SDE Noise)",
|
|
default = 0.0,
|
|
min = 0.0,
|
|
max = 1.0,
|
|
step = 0.05,
|
|
hint = "Ancestral noise injection. 0=deterministic ODE. Scales with t_prev.",
|
|
},
|
|
{
|
|
key = "seed",
|
|
type = "slider",
|
|
label = "Seed",
|
|
default = 42,
|
|
min = 0,
|
|
max = 999999,
|
|
step = 1,
|
|
hint = "RNG seed for SDE noise.",
|
|
},
|
|
},
|
|
}
|
|
|
|
local EPSILON = 1e-8
|
|
|
|
-- Module state -- reset on step_idx_==0
|
|
local _sigma_max = nil
|
|
local _x_prev_lb = nil -- look-back previous x (before update)
|
|
|
|
local function clamp(v, lo, hi)
|
|
if v < lo then return lo end
|
|
if v > hi then return hi end
|
|
return v
|
|
end
|
|
|
|
local function make_rng(seed)
|
|
local state = math.floor(seed) % 2147483647
|
|
if state <= 0 then state = state + 2147483646 end
|
|
return function()
|
|
state = (state * 1664525 + 1013904223) % 2147483648
|
|
return state / 2147483648.0
|
|
end
|
|
end
|
|
|
|
local function normal(u1, u2)
|
|
return math.sqrt(-2.0 * math.log(math.max(u1, EPSILON))) * math.cos(2.0 * math.pi * u2)
|
|
end
|
|
|
|
local function l2_norm(arr, n)
|
|
local s = 0.0
|
|
for i = 0, n - 1 do s = s + arr[i] * arr[i] end
|
|
return math.sqrt(s + EPSILON)
|
|
end
|
|
|
|
-- Shape decomposition: returns unit direction table + L2 norm (mean_scale)
|
|
local function decompose_shape(vt, n)
|
|
local norm = l2_norm(vt, n)
|
|
local inv = norm > EPSILON and (1.0 / norm) or 0.0
|
|
local shape = {}
|
|
for i = 0, n - 1 do shape[i] = vt[i] * inv end
|
|
return shape, norm
|
|
end
|
|
|
|
-- Gram-Schmidt drift guard: projects shape onto orthogonal complement of xt
|
|
local function apply_drift_guard(shape, xt, n, threshold)
|
|
local norm_x = l2_norm(xt, n)
|
|
if norm_x < EPSILON then return shape end
|
|
|
|
local inv_x = 1.0 / norm_x
|
|
local dot_sx = 0.0
|
|
for i = 0, n - 1 do dot_sx = dot_sx + shape[i] * (xt[i] * inv_x) end
|
|
|
|
if math.abs(dot_sx) <= threshold then return shape end
|
|
|
|
local proj = {}
|
|
for i = 0, n - 1 do
|
|
proj[i] = shape[i] - dot_sx * (xt[i] * inv_x)
|
|
end
|
|
|
|
local proj_norm = l2_norm(proj, n)
|
|
if proj_norm < EPSILON then return shape end
|
|
|
|
local inv_proj = 1.0 / proj_norm
|
|
for i = 0, n - 1 do proj[i] = proj[i] * inv_proj end
|
|
return proj
|
|
end
|
|
|
|
-- Look-back SNR smoother: blend x_curr toward x_prev, lambda fades with sigma
|
|
local function look_back_smooth(x_curr, x_prev, t_curr, sigma_max, lb_lambda, snr_power, n)
|
|
if x_prev == nil or lb_lambda < EPSILON then return x_curr, 0.0 end
|
|
local ratio = math.min(t_curr / math.max(sigma_max, EPSILON), 1.0)
|
|
local lam = lb_lambda * (ratio ^ snr_power)
|
|
local out = {}
|
|
for i = 0, n - 1 do
|
|
out[i] = (1.0 - lam) * x_curr[i] + lam * x_prev[i]
|
|
end
|
|
return out, lam
|
|
end
|
|
|
|
-- ── step() function ───────────────────────────────────────────────────────────
|
|
|
|
function step(xt, vt, t_curr, t_prev, n)
|
|
local rw = clamp((params and params.relational_weight) or 0.5, 0.0, 1.0)
|
|
local sig_power = clamp((params and params.sigma_power) or 1.0, 0.25, 4.0)
|
|
local lb_lambda = clamp((params and params.look_back_lambda) or 0.0, 0.0, 0.5)
|
|
local lb_snr_pow = clamp((params and params.look_back_snr_power) or 1.5, 0.5, 3.0)
|
|
local drift_on = (params and params.drift_guard) or false
|
|
local drift_thr = clamp((params and params.drift_threshold) or 0.85, 0.1, 1.0)
|
|
local eta = clamp((params and params.eta) or 0.0, 0.0, 1.0)
|
|
local seed = math.floor((params and params.seed) or 42)
|
|
|
|
local step_idx_ = step_index or 0
|
|
|
|
-- Reset state at start of each generation
|
|
if step_idx_ == 0 then
|
|
_sigma_max = t_curr
|
|
_x_prev_lb = nil
|
|
end
|
|
if _sigma_max == nil then _sigma_max = t_curr end
|
|
|
|
-- Snapshot xt before update for look-back (copy to plain table)
|
|
local x_curr_snapshot = nil
|
|
if lb_lambda > EPSILON then
|
|
x_curr_snapshot = {}
|
|
for i = 0, n - 1 do x_curr_snapshot[i] = xt[i] end
|
|
end
|
|
|
|
-- ── 1. Shape decomposition ────────────────────────────────────────────────
|
|
local shape_vec, mean_scale = decompose_shape(vt, n)
|
|
|
|
-- ── 2. Optional drift guard ───────────────────────────────────────────────
|
|
if drift_on and drift_thr < 1.0 then
|
|
-- xt is a plain table in step() -- pass directly
|
|
local xt_tbl = {}
|
|
for i = 0, n - 1 do xt_tbl[i] = xt[i] end
|
|
shape_vec = apply_drift_guard(shape_vec, xt_tbl, n, drift_thr)
|
|
end
|
|
|
|
-- ── 3. Sigma-adaptive relational weight ───────────────────────────────────
|
|
-- rw_eff = rw * (t_curr / sigma_max) ^ sigma_power
|
|
-- At high sigma: rw_eff = rw (full relational). At low sigma: fades to 0.
|
|
local sigma_ratio = math.min(t_curr / math.max(_sigma_max, EPSILON), 1.0)
|
|
local rw_eff = rw * (sigma_ratio ^ sig_power)
|
|
|
|
-- ── 4. Blend + Euler update ───────────────────────────────────────────────
|
|
local dt = t_prev - t_curr -- negative in flow-matching
|
|
local x_next = {}
|
|
for i = 0, n - 1 do
|
|
local rel_i = shape_vec[i] * mean_scale
|
|
local eff_vt_i = rw_eff * rel_i + (1.0 - rw_eff) * vt[i]
|
|
x_next[i] = xt[i] + dt * eff_vt_i
|
|
end
|
|
|
|
-- ── 5. Look-back smoother ─────────────────────────────────────────────────
|
|
if lb_lambda > EPSILON then
|
|
local lam
|
|
x_next, lam = look_back_smooth(x_next, _x_prev_lb, t_curr, _sigma_max, lb_lambda, lb_snr_pow, n)
|
|
_x_prev_lb = x_curr_snapshot
|
|
end
|
|
|
|
-- ── 6. Write x_next back to xt ───────────────────────────────────────────
|
|
for i = 0, n - 1 do xt[i] = x_next[i] end
|
|
|
|
-- ── 7. Optional SDE noise ─────────────────────────────────────────────────
|
|
if eta > 0.0 and t_prev > EPSILON then
|
|
local rng = make_rng(seed + step_idx_ * 7919)
|
|
local scale = t_prev * eta
|
|
for i = 0, n - 1 do
|
|
local u1 = math.max(rng(), EPSILON)
|
|
local u2 = rng()
|
|
xt[i] = xt[i] + normal(u1, u2) * scale
|
|
end
|
|
end
|
|
end
|