digital-garbageandClaude Opus 5.5
ce2de22b2e
STAS (arXiv 2603.17825) sets massive-activation channels on frame 0 and frame edges to alpha x peak at one block over the first 40% of steps. Rating-learned dials now share rated_dial.py. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
67 lines
2.7 KiB
Python
67 lines
2.7 KiB
Python
"""STAS: steer the massive activations of a video DiT (arXiv 2603.17825).
|
|
|
|
A few channels of a video transformer's hidden state carry huge values ("massive
|
|
activations"), strongest on the first frame and on the edges of each latent frame. They
|
|
act as the model's own rescalers. STAS sets them, at one block and over the first steps
|
|
only, to alpha x their current peak:
|
|
|
|
g(i, d) = alpha * max_j |H[j, d]| * sign(H[i, d]) for i in S, d in M
|
|
|
|
S = every token of the first latent frame + the first and last p% of every frame's
|
|
tokens; M = channels whose peak is over 50x the mean magnitude. Paper settings: one
|
|
block ~30% deep, first 20 of 50 steps, p = 8%, alpha 1.2-2.5 (tested on Wan, CogVideoX).
|
|
No extra model call. Alpha is learned from ratings (rated_dial.py) or set by hand.
|
|
"""
|
|
|
|
import math
|
|
|
|
import torch
|
|
|
|
try:
|
|
from .rated_dial import Dial, MODES # noqa: F401
|
|
except ImportError:
|
|
from rated_dial import Dial, MODES # noqa: F401
|
|
|
|
DIAL = Dial("stas", start=2.0, lo=0.5, hi=3.0, explore=0.2)
|
|
DEFAULT_BLOCK = 15 # ~30% of H3's 50, the paper's depth (block 9 of Wan's 30)
|
|
EARLY_FRACTION = 0.4 # first 20 of 50 steps
|
|
EDGE = 0.08 # p: head and tail share of each frame's tokens
|
|
MA_RATIO = 50.0
|
|
MAX_DIMS = 8
|
|
|
|
|
|
def first_steps(total):
|
|
"""How many of `total` steps are steered."""
|
|
return max(1, math.ceil(EARLY_FRACTION * int(total)))
|
|
|
|
|
|
def target_rows(frames, per_frame, edge=EDGE):
|
|
"""Row indices (within the video rows) of S: all of frame 0, then the head and tail
|
|
`edge` share of every frame. Rows are frame-major, as H3 packs them."""
|
|
k = max(1, round(edge * per_frame))
|
|
within = torch.cat([torch.arange(k), torch.arange(per_frame - k, per_frame)]).unique()
|
|
rows = (torch.arange(frames)[:, None] * per_frame + within[None, :]).flatten()
|
|
return torch.cat([torch.arange(per_frame), rows]).unique()
|
|
|
|
|
|
def ma_dims(video):
|
|
"""-> (channel indices, their peaks, their peak/mean ratios). `video` is [n, D]."""
|
|
peak = torch.maximum(video.amax(dim=0).float().abs(), video.amin(dim=0).float().abs())
|
|
stride = max(1, video.shape[0] // 1024)
|
|
mean = video[::stride].float().abs().mean().clamp(min=1e-12)
|
|
ratio = peak / mean
|
|
dims = torch.nonzero(ratio > MA_RATIO).flatten()
|
|
if dims.numel() > MAX_DIMS:
|
|
dims = dims[ratio[dims].topk(MAX_DIMS).indices]
|
|
return dims, peak[dims], ratio[dims]
|
|
|
|
|
|
def steer(video, rows, dims, peaks, alpha):
|
|
"""Set video[rows, dims] to alpha * peak * sign, in place."""
|
|
if not dims.numel() or not rows.numel():
|
|
return video
|
|
r, d = rows.to(video.device)[:, None], dims.to(video.device)[None, :]
|
|
target = (alpha * peaks.to(video.device))[None, :] * torch.sign(video[r, d].float())
|
|
video[r, d] = target.to(video.dtype)
|
|
return video
|