programdemos/rank-and-scaling/rank-and-scaling.xtl
Rank and scaling: the adapter's rank swept from 1 to 8 on the learning task of lesson 3 (a frozen 8 x 8 layer, examples from a hidden shift of it, 8 trained on and 4 held out), under two rules for the strength s: LoRA's alpha / r (lo:s_caling) and rsLoRA's alpha / sqrt r (lo:s_calingRs). For each rank and rule, after the same number of Adam steps from the same kind of start: the numbers trained (analytical), the strength (analytical), the loss per example on the examples trained on and on the held-out ones (measured), and the size of the adapter's table s B A (measured: the square root of the sum of its squares).
Data
alpha : Float
The settings: alpha, the learning rate, how many of the 12 examples are trained on (the last 4 are always held out), and the steps each rank is trained for.
alpha ← 4.0
i : Float
The task of lesson 3: the frozen layer W by a fixed ripple, the hidden shift T (8 x 8, of rank 4), 12 inputs of 8 numbers as columns with the outputs the shifted layer W + T wants for them, the first n trained on, the last 4 held out.
i ← f̲loat r̲ange 64
T : Float
T ← (8 4 r̲eshape 0.5 × s̲in 1.3 × p × p) '+ '× i̲nner 4 8 r̲eshape 0.5 × c̲os 0.6 × p × p
The sweep
table : Float
Rank 1 to 8, each trained steps times by Adam (lo:a_dam) from the start LoRA uses (B at 0, A seeded by the rank), once under each rule for s. A row per rank: the numbers trained; then under alpha / r: s, the loss per example on the examples trained on, on the held-out ones, and the size of s B A; then the same four under alpha / sqrt r.
(_, table) ← 8 '{ (k, rows) → a1 ← (W, 1 ˡᵒs̲eeded k c̲at 8, (8 c̲at k) r̲eshape 0.0, alpha ˡᵒs̲caling k) a2 ← (W, 1 ˡᵒs̲eeded k c̲at 8, (8 c̲at k) r̲eshape 0.0, alpha ˡᵒs̲calingRs k) (t1, _, _) ← steps '{ e → lr ˡᵒa̲dam (e, (Xt, Yt)) } p̲ower ˡᵒa̲damStart a1 (t2, _, _) ← steps '{ e → lr ˡᵒa̲dam (e, (Xt, Yt)) } p̲ower ˡᵒa̲damStart a2 (_, A1, B1, s1) ← t1 (_, A2, B2, s2) ← t2 d1 ← ˡᵒd̲elta (A1, B1, s1) d2 ← ˡᵒd̲elta (A2, B2, s2) row ← (f̲loat k ˡᵒp̲arams 8 8) c̲at s1 c̲at ((t1 ˡᵒl̲oss (Xt, Yt)) ÷ f̲loat n) c̲at ((t1 ˡᵒl̲oss (Xh, Yh)) ÷ 4.0) c̲at (s̲qrt '+ r̲/ r̲avel d1 × d1) c̲at s2 c̲at ((t2 ˡᵒl̲oss (Xt, Yt)) ÷ f̲loat n) c̲at ((t2 ˡᵒl̲oss (Xh, Yh)) ÷ 4.0) c̲at s̲qrt '+ r̲/ r̲avel d2 × d2 (k + 1, rows c̲at row) } p̲ower (1, 0 r̲eshape 0.0)
_ : Char
The measurements as three charts, rank 1 to 8 along the bottom: the two losses under alpha / r, the same under alpha / sqrt r, and the size of s B A under each rule beside the size of T.
_ ← ("Loss by rank, s = alpha / r" "rank" "loss per example" "trained on" "held out") ᵖc̲hart! '{ j → j s̲elect₂ M } m̲ap 3 4