42 lines
1.5 KiB
Lua
42 lines
1.5 KiB
Lua
-- rescaled_cfg.lua: Std-matched post-processing guidance
|
|
-- Runs APG at full scale, then rescales output to match conditional std.
|
|
-- Routes through native APG (momentum + projection + norm thresholding),
|
|
-- then applies std-matching post-processing.
|
|
|
|
guidance = {
|
|
name = "rescaled_cfg",
|
|
display = "Rescaled CFG",
|
|
description = "Std-matched to prevent saturation",
|
|
}
|
|
|
|
function guide(pred_cond, pred_uncond, guidance_scale, result, Oc, T, norm_threshold)
|
|
local n = Oc * T
|
|
local phi = (guidance_scale > 4.0) and 0.95 or 0.7
|
|
|
|
-- Run APG at full guidance scale first
|
|
apg(pred_cond, pred_uncond, guidance_scale, result, Oc, T, norm_threshold)
|
|
|
|
-- Compute std of conditional prediction and guided output
|
|
local sum_c, sum2_c = 0, 0
|
|
local sum_g, sum2_g = 0, 0
|
|
for i = 0, n - 1 do
|
|
local c = pred_cond[i]
|
|
local g = result[i]
|
|
sum_c = sum_c + c; sum2_c = sum2_c + c * c
|
|
sum_g = sum_g + g; sum2_g = sum2_g + g * g
|
|
end
|
|
local mean_c = sum_c / n
|
|
local mean_g = sum_g / n
|
|
local var_c = sum2_c / n - mean_c * mean_c
|
|
local var_g = sum2_g / n - mean_g * mean_g
|
|
local std_c = (var_c > 0) and math.sqrt(var_c) or 1e-5
|
|
local std_g = (var_g > 0) and math.sqrt(var_g) or 1e-5
|
|
|
|
-- Rescale to match conditional std, blend with raw APG output
|
|
local factor = std_c / (std_g + 1e-5)
|
|
for i = 0, n - 1 do
|
|
local rescaled = result[i] * factor
|
|
result[i] = phi * rescaled + (1 - phi) * result[i]
|
|
end
|
|
end
|