mirror of
https://github.com/wassname/lora-lite.git
synced 2026-08-11 11:21:06 +08:00
variants: replace arrow's dense block with diagonal-plus-low-rank core
antipasto_arrow -> antipasto_dplr. The arrowhead's dense b x b block is the wrong
shape: b^2 params, mixes only the top-b, and sits on the S-scaled coords so its
perturbation is amplified by the largest singular values (block=128 collapsed to
45.7% at the gain's lr). Replace it with LoRA's lesson -- a low-rank core inside
the frozen basis, ADDED to the gain:
DeltaW = U [diag(S_eff) + coeff * B A] Vh, A:(k,r) B:(r,k), B=0 at init
The low-rank part mixes the whole top-r subspace for 2*r*k params (k=LoRA's rank),
and being additive (not * diag(S)) it is S-independent -- the amplification edge is
gone by construction. Diagonal gain unchanged; identity at init from B=0 and g=0.
Wired through benchmark (antipasto_lora_rank, run_id __k suffix), justfile, cost_report,
smoke (green, dplr attaches/trains/round-trips). Arrow code removed; its run results
stay on disk for comparison.
Co-Authored-By: Claudypoo <noreply@anthropic.com>
This commit is contained in:
wassname
@@ -23,7 +23,7 @@ from .variants.antipasto import AntiPaSTOConfig
|
||||
from .variants.antipasto_rot import AntiPaSTORotConfig
|
||||
from .variants.antipasto_ablate import AntiPaSTOAblateConfig
|
||||
from .variants.antipasto_corda import AntiPaSTOCorDAConfig
|
||||
from .variants.antipasto_arrow import AntiPaSTOArrowConfig
|
||||
from .variants.antipasto_dplr import AntiPaSTODPLRConfig
|
||||
from .variants.road import RoadConfig
|
||||
|
||||
__all__ = [
|
||||
@@ -40,7 +40,7 @@ __all__ = [
|
||||
"AntiPaSTORotConfig",
|
||||
"AntiPaSTOAblateConfig",
|
||||
"AntiPaSTOCorDAConfig",
|
||||
"AntiPaSTOArrowConfig",
|
||||
"AntiPaSTODPLRConfig",
|
||||
"RoadConfig",
|
||||
"attach",
|
||||
"detach",
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
from . import ( # noqa: F401 side-effect: register
|
||||
lora, pissa, delora, ia3, dora, hra, eva, antipasto, road,
|
||||
antipasto_rot, antipasto_ablate, antipasto_corda, antipasto_arrow,
|
||||
antipasto_rot, antipasto_ablate, antipasto_corda, antipasto_dplr,
|
||||
)
|
||||
|
||||
+61
-53
@@ -1,22 +1,33 @@
|
||||
"""AntiPaSTO-Arrow: cross-direction mixing via a cheap arrowhead core.
|
||||
"""AntiPaSTO-DPLR: diagonal-plus-low-rank core in the frozen SVD basis.
|
||||
|
||||
antipasto's core is diagonal (S_eff = S * gain): it reweights each singular direction
|
||||
independently but cannot let direction i drive direction j. A full dense r x r core
|
||||
restores all mixing but costs r^2 params. The arrowhead is the cheap middle: a dense
|
||||
block on the top-b directions (where the action lives), the diagonal gain on the rest.
|
||||
antipasto's core is diagonal (a per-direction gain); it rescales each singular
|
||||
direction but cannot mix one into another. The arrowhead tried a dense b x b block
|
||||
on the top-b directions, but a dense block is the wrong shape (b^2 params, mixes only
|
||||
the top-b) and -- sitting on the S-scaled coords -- its perturbation is amplified by
|
||||
the largest singular values, so it destabilizes. The fix is LoRA's lesson: a low-rank
|
||||
core. Put a trainable rank-k core inside the frozen U/Vh basis, ADDED to the gain:
|
||||
|
||||
core C (r x r, on the S-scaled coords):
|
||||
[ B (b x b dense) | 0 ] B = I_b + coeff*M (top-b mixing)
|
||||
[ 0 | diag(1 + ELU(coeff*g)) ] tail = antipasto's gain
|
||||
DeltaW = U @ C @ diag(S) @ Vh
|
||||
cost: b^2 + (r-b) params, one b x b matmul per forward.
|
||||
W = U diag(S) Vh + W_res # frozen top-r SVD
|
||||
learn: g (r,) # diagonal gain
|
||||
A (k,r), B (r,k) # low-rank mixing core, B=0 at init
|
||||
S_eff = S * (1 + ELU(coeff * g))
|
||||
y = x @ W_res.T + ( (Vh x) * S_eff + coeff * B (A (Vh x)) ) @ U.T
|
||||
|
||||
Identity at init: M=0 -> B=I, g=0 -> 1+ELU(0)=1, so C=I and DeltaW = U diag(S) Vh.
|
||||
coeff=0 -> C=I too (runtime off). The block is the linear (1+z) regime -- stable but
|
||||
not strictly bounded; for a can't-blow-up guarantee on the top directions use
|
||||
antipasto_ablate.
|
||||
so the trainable core is C = diag(S_eff) + coeff * B A acting in S-space, and
|
||||
DeltaW = U C Vh. The diagonal part scales directions; the low-rank part B A mixes them
|
||||
across the whole top-r subspace for 2*r*k params (k=LoRA's rank), not b^2.
|
||||
|
||||
Refs: antipasto.py (diagonal sibling), antipasto_corda.py (off-axis basis argument).
|
||||
Why the low-rank part is ADDED, not multiplied into diag(S): an additive core
|
||||
U (BA) Vh is independent of S, so a unit step in BA moves W by O(1), not O(S). That is
|
||||
exactly the S-amplification edge that made the dense arrowhead block blow up at the
|
||||
gain's learning rate -- gone by construction.
|
||||
|
||||
Identity at init: B=0 -> BA=0, g=0 -> 1+ELU(0)=1, so C=diag(S) and DeltaW = U diag(S) Vh.
|
||||
coeff=0 -> identity too (runtime off). The basis (U, Vh) stays frozen and interpretable;
|
||||
only the gain and the rank-k core move.
|
||||
|
||||
Refs: antipasto.py (diagonal sibling), lora.py (the low-rank core), antipasto_corda.py
|
||||
(oriented basis -- composes with this core).
|
||||
"""
|
||||
from dataclasses import dataclass
|
||||
from typing import Iterable, Literal
|
||||
@@ -36,42 +47,42 @@ CalibrationData = Iterable[CalibrationBatch]
|
||||
|
||||
@register_config
|
||||
@dataclass
|
||||
class AntiPaSTOArrowConfig(AdapterConfig):
|
||||
variant: str = "antipasto_arrow"
|
||||
class AntiPaSTODPLRConfig(AdapterConfig):
|
||||
variant: str = "antipasto_dplr"
|
||||
r: int = 256
|
||||
# Dense interaction block on the top-b singular directions; sets capacity and the
|
||||
# only quadratic cost (b^2 params/module). b=1 degenerates to antipasto; b->r
|
||||
# approaches a full dense r-core (~LoRA params) at the cost arrow exists to avoid.
|
||||
block: int = 8
|
||||
suppress_only: bool = False # clamp the tail g<=0 (attenuate only); block unaffected.
|
||||
# Tail guarantee holds for coeff>=0; coeff<0 inverts the product and re-amplifies.
|
||||
coeff: float = 1.0 # runtime knob: 0=identity, scales both block and tail
|
||||
act_pool: Literal["rms", "mean_abs"] = "rms" # group_init selection, see antipasto
|
||||
# Rank of the low-rank mixing core (LoRA's r, but inside the frozen subspace).
|
||||
# Params = r (gain) + 2*r*lora_rank. k=0 degenerates to plain antipasto.
|
||||
lora_rank: int = 8
|
||||
suppress_only: bool = False # clamp the gain g<=0 (attenuate only); core unaffected.
|
||||
coeff: float = 1.0 # runtime knob: 0=identity, scales gain and core.
|
||||
act_pool: Literal["rms", "mean_abs"] = "rms" # group_init selection, see antipasto.
|
||||
|
||||
|
||||
@register
|
||||
class AntiPaSTOArrow:
|
||||
name = "antipasto_arrow"
|
||||
class AntiPaSTODPLR:
|
||||
name = "antipasto_dplr"
|
||||
|
||||
@staticmethod
|
||||
def param_specs(d_in, d_out, cfg):
|
||||
r, b = cfg.r, cfg.block
|
||||
if not 1 <= b < r:
|
||||
raise ValueError(f"antipasto_arrow needs 1 <= block({b}) < r({r}).")
|
||||
r, k = cfg.r, cfg.lora_rank
|
||||
if not 0 < k <= r:
|
||||
raise ValueError(f"antipasto_dplr needs 0 < lora_rank({k}) <= r({r}).")
|
||||
return dict(
|
||||
lora_U=ParamSpec((d_out, r), init="zeros", trainable=False, as_buffer=True),
|
||||
lora_S=ParamSpec((r,), init="zeros", trainable=False, as_buffer=True),
|
||||
lora_Vh=ParamSpec((r, d_in), init="zeros", trainable=False, as_buffer=True),
|
||||
# Dense b x b interaction on the top-b directions. init 0 -> B=I -> identity.
|
||||
lora_M=ParamSpec((b, b), init="zeros"),
|
||||
# Diagonal bounded gain on the remaining r-b directions (== antipasto's g).
|
||||
lora_g=ParamSpec((r - b,), init="zeros"),
|
||||
# Diagonal gain (== antipasto). init 0 -> 1+ELU(0)=1 -> identity.
|
||||
lora_g=ParamSpec((r,), init="zeros"),
|
||||
# Low-rank core B@A in the frozen subspace. A down (r->k), B up (k->r).
|
||||
# B=0 at init -> core=0 -> identity (LoRA convention).
|
||||
lora_A=ParamSpec((k, r), init="kaiming"),
|
||||
lora_B=ParamSpec((r, k), init="zeros"),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def init(layer: nn.Module, cfg) -> None:
|
||||
if type(layer) is not nn.Linear:
|
||||
raise TypeError("AntiPaSTOArrow mutates layer.weight into W_res; nn.Linear only.")
|
||||
raise TypeError("AntiPaSTODPLR mutates layer.weight into W_res; nn.Linear only.")
|
||||
with torch.no_grad():
|
||||
W = layer.weight.data.float()
|
||||
U, S, Vh = torch.linalg.svd(W, full_matrices=False)
|
||||
@@ -85,9 +96,9 @@ class AntiPaSTOArrow:
|
||||
|
||||
@staticmethod
|
||||
def group_init(model: nn.Module, targets, cfg, calibration_data: CalibrationData | None) -> None:
|
||||
"""Wanda-style data-driven dimension selection, identical to antipasto: re-pick
|
||||
the top-r directions by S[i] * pool|X @ Vh[i]|. Runs before training (g, M at
|
||||
their zero init), so re-selecting the basis is a harmless no-op on the core."""
|
||||
"""Wanda-style re-selection of the top-r directions, identical to antipasto.
|
||||
Runs before training while g and B are still zero, so the core contributes
|
||||
nothing and re-selecting the basis is a no-op on the adapter output."""
|
||||
if calibration_data is None:
|
||||
return
|
||||
|
||||
@@ -122,15 +133,13 @@ class AntiPaSTOArrow:
|
||||
for name, layer in layers.items():
|
||||
X = torch.cat(captured[name], dim=0)
|
||||
if X.shape[0] < r:
|
||||
raise RuntimeError(f"AntiPaSTOArrow at {name}: {X.shape[0]} tokens, need >= r={r}")
|
||||
raise RuntimeError(f"AntiPaSTODPLR at {name}: {X.shape[0]} tokens, need >= r={r}")
|
||||
# Rebuild the FULL W exactly (W_res + stored top-r), then re-select top-r.
|
||||
W_res = layer.weight.data.float()
|
||||
W_orig = W_res + (layer.lora_U.float() * layer.lora_S.float()) @ layer.lora_Vh.float()
|
||||
U_full, S_full, Vh_full = torch.linalg.svd(W_orig, full_matrices=False)
|
||||
proj = X.to(Vh_full) @ Vh_full.T
|
||||
act_mag = proj.pow(2).mean(0).sqrt() if pool == "rms" else proj.abs().mean(0)
|
||||
# Pick top-r by score, then sort by SVD index. svd() returns S descending,
|
||||
# so the block's first-b coords are the b largest-S among the selected r
|
||||
# (= where the action lives), not the b highest-score.
|
||||
idx = (S_full * act_mag).argsort(descending=True)[:r].sort().values
|
||||
Ur, Sr, Vhr = U_full[:, idx], S_full[idx], Vh_full[idx]
|
||||
W_res_new = (W_orig - (Ur * Sr) @ Vhr).to(layer.weight.dtype)
|
||||
@@ -150,19 +159,18 @@ class AntiPaSTOArrow:
|
||||
U = layer.lora_U.to(x.dtype) # (d_out, r)
|
||||
S = layer.lora_S.to(x.dtype) # (r,)
|
||||
Vh = layer.lora_Vh.to(x.dtype) # (r, d_in)
|
||||
M = layer.lora_M.to(x.dtype) # (b, b)
|
||||
g = layer.lora_g.to(x.dtype) # (r-b,)
|
||||
coeff, b = float(cfg.coeff), cfg.block
|
||||
g = layer.lora_g.to(x.dtype) # (r,)
|
||||
A = layer.lora_A.to(x.dtype) # (k, r)
|
||||
B = layer.lora_B.to(x.dtype) # (r, k)
|
||||
coeff = float(cfg.coeff)
|
||||
|
||||
cS = (x @ Vh.T) * S # (..., r) = diag(S) Vh x
|
||||
|
||||
# Top-b: dense block B = I_b + coeff*M couples the top singular directions.
|
||||
eye = torch.eye(b, dtype=x.dtype, device=x.device)
|
||||
top = cS[..., :b] @ (eye + coeff * M).T # (..., b)
|
||||
# Tail: antipasto's bounded diagonal gain (see antipasto.py for the 1+ELU why).
|
||||
if cfg.suppress_only:
|
||||
g = torch.clamp(g, max=0.0)
|
||||
tail = cS[..., b:] * (1.0 + F.elu(coeff * g)) # (..., r-b)
|
||||
|
||||
h = torch.cat([top, tail], dim=-1) # (..., r)
|
||||
p = x @ Vh.T # (..., r) = Vh x (unscaled)
|
||||
S_eff = S * (1.0 + F.elu(coeff * g)) # diagonal gain (see antipasto.py)
|
||||
# Diagonal part scales each direction; low-rank part B@A mixes across the
|
||||
# subspace. Additive (not * diag(S)), so the core is S-independent: a unit
|
||||
# step in B@A moves W by O(1), not O(S) -- no S-amplification edge.
|
||||
h = p * S_eff + coeff * (p @ A.T) @ B.T # (..., r)
|
||||
return y + h @ U.T
|
||||
Reference in New Issue
Block a user