fix: H3 attention decoupling was scaling K magnitude, which barely moves softmax
Live 0-1 strength sweep came back flat -- confirmed the mechanism, not just the value, was wrong. H3's Q/K get RMSNorm'd before attention, so scaling a key toward zero pushes its logit toward the *mean* of an already near-zero-centered distribution, not toward excluded. Switched to what actually suppresses a softmax entry: an additive negative bias on the targeted logits, still passed through the same per-query-group split (still no S^2 mask, still same total FLOPs -- the bias is a per-key vector broadcast over the group's queries). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Sonnet 5
parent
9bbdfc8b4b
commit
1e4553e658
@@ -336,7 +336,7 @@
|
||||
detail: "Blocks 40-49 carry the most residual movement (48+49 alone about a third of it), so a steering push has the most to act on there. Blocks 0-1 are structurally distinct rather than just 'quiet early blocks' — 0 is the raw input projection, 1 does the first cross-patch mixing — expect steering there to behave differently in kind, not just 'weaker', than the middle of the stack." },
|
||||
{ name: "h3_av_decouple", label: "Video/audio attention decoupling (H3, experimental)", kind: "float", default: 0.0, min: 0.0, max: 1.0, step: 0.01,
|
||||
hint: "EXPERIMENTAL, unvalidated, H3 only. Weakens how much video and audio attend to each other, for the failure pattern where one comes out clean and the other doesn't. 0 = off. Not rating-driven — a manual dial, same as Repeat block(s).",
|
||||
detail: "H3 has no separate cross-attention to mask — video, text and audio share one joint self-attention. Instead of building a mask over that whole sequence (unaffordable at H3's token count), this runs the same attention as separate video-query and audio-query passes against a version of the keys where the other modality is scaled down by this amount. Same total cost as the ordinary single pass. 1.0 cuts cross-modal attention entirely, which likely also kills legitimate audio/video sync — search down from a small value, not up to it." },
|
||||
detail: "H3 has no separate cross-attention to mask — video, text and audio share one joint self-attention. Instead of building a mask over that whole sequence (unaffordable at H3's token count), this runs the same attention as separate video-query and audio-query passes, each carrying a penalty on the other modality's attention weight. Same total cost as the ordinary single pass. 1.0 cuts cross-modal attention entirely, which likely also kills legitimate audio/video sync — search down from a small value, not up to it. (v2: the first version tried this by shrinking the other modality's keys instead of penalizing attention directly, and live testing showed almost no effect across the whole 0-1 range — fixed to actually penalize.)" },
|
||||
{ name: "identity_transfer_enabled", label: "Best-FaceID compatibility", kind: "bool", default: false,
|
||||
hint: "Feeds the identity pin image the way Best-FaceID identity LoRAs expect it. Needs an Identity pin set.",
|
||||
detail: "Replaces Continuity's identity-pin guide with separate, non-rendered reference tokens plus an optional ArcFace projector below. Load the LoRA itself in Models. No effect without an identity pin." },
|
||||
|
||||
+33
-14
@@ -2808,7 +2808,7 @@ class FunPackLTXAVSceneChainSampler:
|
||||
}),
|
||||
"h3_av_decouple": ("FLOAT", {
|
||||
"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01,
|
||||
"tooltip": "EXPERIMENTAL, unvalidated (H3 only). Damps how much video and audio attend to EACH OTHER inside H3's joint self-attention, without touching any hidden state -- for the failure pattern where one comes out clean and the other doesn't. 0 = off (installs nothing, zero cost). Higher = weaker cross-modal influence; 1.0 fully cuts it, which likely also kills legitimate AV sync, so this is a knob to search down from, not up to.",
|
||||
"tooltip": "EXPERIMENTAL, unvalidated (H3 only). Damps how much video and audio attend to EACH OTHER inside H3's joint self-attention, without touching any hidden state -- for the failure pattern where one comes out clean and the other doesn't. 0 = off (installs nothing, zero cost). Higher = weaker cross-modal influence via a negative bias on the other modality's attention logits (v2: the first version scaled key magnitude instead and measured flat 0-1 live -- fixed). 1.0 fully cuts it, which likely also kills legitimate AV sync, so this is a knob to search down from, not up to.",
|
||||
}),
|
||||
"alg_guide_blur_strength": ("FLOAT", {
|
||||
"default": 2.0, "min": 1.0, "max": 4.0, "step": 0.1,
|
||||
@@ -6831,12 +6831,20 @@ class FunPackLTXAVSceneChainSampler:
|
||||
sequence). This does not build that S×S mask at all. It uses comfy's own
|
||||
`optimized_attention_override` extension point (comfy/ldm/modules/attention.py's
|
||||
`wrap_attn` -- already wired into every backend, nothing to monkeypatch) to run the
|
||||
SAME joint attention as separate video-query and audio-query passes against a K
|
||||
where the OTHER modality's rows are scaled down by `strength`; text/ref/cond rows
|
||||
SAME joint attention as separate video-query and audio-query passes, each carrying
|
||||
an additive negative bias on the OTHER modality's logits (a per-key vector, not a
|
||||
[seq,seq] matrix -- broadcasts over every query in the group); text/ref/cond rows
|
||||
stay on one ordinary pass, untouched. Splitting the query dimension like this costs
|
||||
nothing extra -- same total Q*K*V work as the single joint call, just partitioned --
|
||||
and the only new memory is two full-length clones of K (audio and video zeroed out
|
||||
by contrast, at most `strength`), not a dense attention matrix.
|
||||
nothing extra -- same total Q*K*V work as the single joint call, just partitioned.
|
||||
|
||||
First cut of this scaled K's magnitude toward zero instead of biasing the logit --
|
||||
measured flat live across the whole 0-1 strength range. After H3's own per-head
|
||||
RMSNorm on Q/K, attention logits already sit close to 0, so shrinking a key's
|
||||
magnitude pushes its logit toward "average", not toward "excluded" -- softmax barely
|
||||
notices. An additive penalty actually pushes the targeted logit negative, which is
|
||||
what exclusion requires. Still unvalidated whether the *hypothesis* (cross-modal
|
||||
attention causes the clean-video/corrupted-audio split) is even correct -- this only
|
||||
fixes the mechanism not doing what the strength dial implied.
|
||||
|
||||
`strength` in [0, 1]; 0 or blank installs nothing (exact pass-through, not just a
|
||||
no-op override) so an unset knob costs literally zero.
|
||||
@@ -6876,6 +6884,14 @@ class FunPackLTXAVSceneChainSampler:
|
||||
to["patches_replace"] = patches_replace
|
||||
|
||||
_mask_cache = {} # seq_len -> (video_mask, audio_mask); one live resolution per gen
|
||||
# Logit bias scale, not a key-magnitude scale (see the method docstring for why
|
||||
# the first cut of this used K-scaling and measured flat across the whole 0-1
|
||||
# range live). After H3's own per-head RMSNorm on Q/K, attention logits already
|
||||
# sit close to 0 -- an additive penalty pushes a targeted key toward genuinely
|
||||
# excluded, where a multiplicative one only pushes it toward "average".
|
||||
# ponytail: fixed heuristic scale, unvalidated against H3's real logit spread --
|
||||
# revisit if the live strength sweep is still flat, or saturates immediately.
|
||||
_BIAS_SCALE = 12.0
|
||||
|
||||
def _override(func, q, k, v, heads, mask=None, skip_reshape=True,
|
||||
transformer_options=None, **kwargs):
|
||||
@@ -6892,18 +6908,21 @@ class FunPackLTXAVSceneChainSampler:
|
||||
if vmask is None or amask is None or k.shape[-2] != seq_len:
|
||||
return func(q, k, v, heads, mask=mask, skip_reshape=skip_reshape, **kwargs)
|
||||
other = ~(vmask | amask)
|
||||
bias_val = -_strength * _BIAS_SCALE
|
||||
out = torch.empty(q.shape[0], seq_len, heads * q.shape[-1],
|
||||
dtype=q.dtype, device=q.device)
|
||||
if amask.any():
|
||||
k_damped = k.clone()
|
||||
k_damped[:, :, vmask, :] *= (1.0 - _strength)
|
||||
out[:, amask, :] = func(q[:, :, amask, :], k_damped, v, heads,
|
||||
mask=None, skip_reshape=skip_reshape, **kwargs)
|
||||
bias = torch.zeros(seq_len, dtype=q.dtype, device=q.device)
|
||||
bias[vmask] = bias_val
|
||||
out[:, amask, :] = func(q[:, :, amask, :], k, v, heads,
|
||||
mask=bias.view(1, 1, 1, seq_len),
|
||||
skip_reshape=skip_reshape, **kwargs)
|
||||
if vmask.any():
|
||||
k_damped = k.clone()
|
||||
k_damped[:, :, amask, :] *= (1.0 - _strength)
|
||||
out[:, vmask, :] = func(q[:, :, vmask, :], k_damped, v, heads,
|
||||
mask=None, skip_reshape=skip_reshape, **kwargs)
|
||||
bias = torch.zeros(seq_len, dtype=q.dtype, device=q.device)
|
||||
bias[amask] = bias_val
|
||||
out[:, vmask, :] = func(q[:, :, vmask, :], k, v, heads,
|
||||
mask=bias.view(1, 1, 1, seq_len),
|
||||
skip_reshape=skip_reshape, **kwargs)
|
||||
if other.any():
|
||||
out[:, other, :] = func(q[:, :, other, :], k, v, heads,
|
||||
mask=None, skip_reshape=skip_reshape, **kwargs)
|
||||
|
||||
@@ -51,8 +51,11 @@ def _qkv(seq_len=6, heads=1, head_dim=2):
|
||||
|
||||
def _fake_func_factory(calls):
|
||||
def fake_func(q, k, v, heads, mask=None, skip_reshape=True, **kw):
|
||||
calls.append({"q_len": q.shape[-2], "k_sum": float(k.sum()), "mask": mask})
|
||||
return torch.full((q.shape[0], q.shape[-2], heads * q.shape[-1]), float(k.sum()))
|
||||
# value encodes the bias so assertions can tell which group produced which output
|
||||
# without needing a real attention computation.
|
||||
marker = float(mask.sum()) if mask is not None else 0.0
|
||||
calls.append({"q_len": q.shape[-2], "mask": mask, "marker": marker})
|
||||
return torch.full((q.shape[0], q.shape[-2], heads * q.shape[-1]), marker)
|
||||
return fake_func
|
||||
|
||||
|
||||
@@ -80,8 +83,8 @@ def test_before_mod_segments_captured_override_is_a_passthrough():
|
||||
q, k, v = _qkv()
|
||||
out = override(_fake_func_factory(calls), q, k, v, heads=1, mask=None, skip_reshape=True)
|
||||
assert len(calls) == 1, "no mod_segments yet -> single unmodified call, not split"
|
||||
assert calls[0]["k_sum"] == float(k.sum())
|
||||
assert torch.allclose(out, torch.full_like(out, float(k.sum())))
|
||||
assert calls[0]["mask"] is None
|
||||
assert torch.allclose(out, torch.zeros_like(out))
|
||||
|
||||
|
||||
def test_an_explicit_mask_is_never_overridden():
|
||||
@@ -96,7 +99,7 @@ def test_an_explicit_mask_is_never_overridden():
|
||||
assert len(calls) == 1 and calls[0]["mask"] is real_mask
|
||||
|
||||
|
||||
def test_splits_into_three_groups_with_correctly_damped_keys():
|
||||
def test_splits_into_three_groups_with_a_bias_penalizing_the_other_modality():
|
||||
node = S()
|
||||
patched = node._install_h3_av_decouple(_FakeModel(), 0.5)
|
||||
_fire_capture_hook(patched)
|
||||
@@ -109,20 +112,27 @@ def test_splits_into_three_groups_with_correctly_damped_keys():
|
||||
by_len = {c["q_len"]: c for c in calls}
|
||||
assert set(by_len) == {len(_VIDEO_IDX), len(_AUDIO_IDX), len(_OTHER_IDX)}
|
||||
|
||||
full_k_sum = float(k.sum()) # 6 rows * 2 dims * 2.0 = 24
|
||||
# audio-query pass: video rows (3 of them) halved in K -> k_sum drops by half of their share
|
||||
per_row = 2 * 2.0 # head_dim * key value
|
||||
audio_call = by_len[len(_AUDIO_IDX)]
|
||||
assert audio_call["k_sum"] == full_k_sum - len(_VIDEO_IDX) * per_row * 0.5
|
||||
video_call = by_len[len(_VIDEO_IDX)]
|
||||
assert video_call["k_sum"] == full_k_sum - len(_AUDIO_IDX) * per_row * 0.5
|
||||
other_call = by_len[len(_OTHER_IDX)]
|
||||
assert other_call["k_sum"] == full_k_sum, "text/other rows ride the ordinary, undamped K"
|
||||
bias_val = -0.5 * 12.0 # strength * _BIAS_SCALE
|
||||
|
||||
# scatter landed each group's output back at its own rows
|
||||
assert torch.allclose(out[0, _AUDIO_IDX], torch.full((len(_AUDIO_IDX), 2), audio_call["k_sum"]))
|
||||
assert torch.allclose(out[0, _VIDEO_IDX], torch.full((len(_VIDEO_IDX), 2), video_call["k_sum"]))
|
||||
assert torch.allclose(out[0, _OTHER_IDX], torch.full((len(_OTHER_IDX), 2), other_call["k_sum"]))
|
||||
audio_call = by_len[len(_AUDIO_IDX)]
|
||||
audio_bias = audio_call["mask"].view(-1)
|
||||
assert torch.allclose(audio_bias[_VIDEO_IDX], torch.full((len(_VIDEO_IDX),), bias_val))
|
||||
assert torch.allclose(audio_bias[_AUDIO_IDX + _OTHER_IDX],
|
||||
torch.zeros(len(_AUDIO_IDX) + len(_OTHER_IDX)))
|
||||
|
||||
video_call = by_len[len(_VIDEO_IDX)]
|
||||
video_bias = video_call["mask"].view(-1)
|
||||
assert torch.allclose(video_bias[_AUDIO_IDX], torch.full((len(_AUDIO_IDX),), bias_val))
|
||||
assert torch.allclose(video_bias[_VIDEO_IDX + _OTHER_IDX],
|
||||
torch.zeros(len(_VIDEO_IDX) + len(_OTHER_IDX)))
|
||||
|
||||
other_call = by_len[len(_OTHER_IDX)]
|
||||
assert other_call["mask"] is None, "text/other rows ride the ordinary, unbiased pass"
|
||||
|
||||
# scatter landed each group's output back at its own rows (marker = mask.sum())
|
||||
assert torch.allclose(out[0, _AUDIO_IDX], torch.full((len(_AUDIO_IDX), 2), audio_call["marker"]))
|
||||
assert torch.allclose(out[0, _VIDEO_IDX], torch.full((len(_VIDEO_IDX), 2), video_call["marker"]))
|
||||
assert torch.allclose(out[0, _OTHER_IDX], torch.zeros(len(_OTHER_IDX), 2))
|
||||
|
||||
|
||||
def test_chains_through_an_already_installed_block0_hook():
|
||||
|
||||
Reference in New Issue
Block a user