rank-and-scaling.xtl

1⍝!/usr/bin/env xetal
2⍝│ Rank and scaling: the adapter's rank swept from 1 to 8 on the
3⍝│ learning task of lesson 3 (a frozen 8 x 8 layer, examples from a
4⍝│ hidden shift of it, 8 trained on and 4 held out), under two rules
5⍝│ for the strength s: LoRA's alpha / r (lo:s_caling) and rsLoRA's
6⍝│ alpha / sqrt r (lo:s_calingRs). For each rank and rule, after the
7⍝│ same number of Adam steps from the same kind of start: the numbers
8⍝│ trained (analytical), the strength (analytical), the loss per
9⍝│ example on the examples trained on and on the held-out ones
10⍝│ (measured), and the size of the adapter's table s B A (measured:
11⍝│ the square root of the sum of its squares).
12
13"lo:" u̲se< "LoRA"
14"p:" u̲se< "Plot"
15
16⍝━━ Data
17
18⍝│ The settings: alpha, the learning rate, how many of the 12 examples
19⍝│ are trained on (the last 4 are always held out), and the steps each
20⍝│ rank is trained for.
21alpha ← 4.0
22lr ← 0.05
23n ← 8
24steps ← 100
25⍝│ The task of lesson 3: the frozen layer W by a fixed ripple, the
26⍝│ hidden shift T (8 x 8, of rank 4), 12 inputs of 8 numbers as
27⍝│ columns with the outputs the shifted layer W + T wants for them,
28⍝│ the first n trained on, the last 4 held out.
29i ← f̲loat r̲ange 64
30W ← 8 8 r̲eshape s̲in 0.7 × i × i
31p ← f̲loat r̲ange 32
32T ← (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
33c ← f̲loat r̲ange 96
34X ← 8 12 r̲eshape c̲os 0.4 × c + 0.05 × c × c
35Y ← (W + T) '+ '× i̲nner X
36Xt ← n t̲ake₂ X
37Yt ← n t̲ake₂ Y
38Xh ← -4 t̲ake₂ X
39Yh ← -4 t̲ake₂ Y
40
41⍝━━ The sweep
42
43⍝│ Rank 1 to 8, each trained steps times by Adam (lo:a_dam) from the
44⍝│ start LoRA uses (B at 0, A seeded by the rank), once under each
45⍝│ rule for s. A row per rank: the numbers trained; then under alpha /
46⍝│ r: s, the loss per example on the examples trained on, on the
47⍝│ held-out ones, and the size of s B A; then the same four under
48⍝│ alpha / sqrt r.
49(_, table) ← 8 '{ (k, rows) →
50  a1 ← (W, 1 ˡᵒs̲eeded k c̲at 8, (8 c̲at k) r̲eshape 0.0, alpha ˡᵒs̲caling k)
51  a2 ← (W, 1 ˡᵒs̲eeded k c̲at 8, (8 c̲at k) r̲eshape 0.0, alpha ˡᵒs̲calingRs k)
52  (t1, _, _) ← steps '{ e → lr ˡᵒa̲dam (e, (Xt, Yt)) } p̲ower ˡᵒa̲damStart a1
53  (t2, _, _) ← steps '{ e → lr ˡᵒa̲dam (e, (Xt, Yt)) } p̲ower ˡᵒa̲damStart a2
54  (_, A1, B1, s1) ← t1
55  (_, A2, B2, s2) ← t2
56  d1 ← ˡᵒd̲elta (A1, B1, s1)
57  d2 ← ˡᵒd̲elta (A2, B2, s2)
58  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
59  (k + 1, rows c̲at row)
60} p̲ower (1, 0 r̲eshape 0.0)
61M ← 8 9 r̲eshape table
62M
63⍝│ The size of the shift T itself (the square root of the sum of its
64⍝│ squares), which a fitted adapter's table comes to, and the loss per
65⍝│ example before any training (B at 0: the layer is W's), on the
66⍝│ examples trained on and on the held-out ones, the same at every rank.
67s̲qrt '+ r̲/ r̲avel T × T
68w0 ← (W, 1 8 r̲eshape 0.0, 8 1 r̲eshape 0.0, 1.0)
69((w0 ˡᵒl̲oss (Xt, Yt)) ÷ f̲loat n) c̲at (w0 ˡᵒl̲oss (Xh, Yh)) ÷ 4.0
70⍝│ The measurements as three charts, rank 1 to 8 along the bottom: the
71⍝│ two losses under alpha / r, the same under alpha / sqrt r, and the
72⍝│ size of s B A under each rule beside the size of T.
73_ ← ("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
74_ ← ("Loss by rank, s = alpha / sqrt r" "rank" "loss per example" "trained on" "held out") ᵖc̲hart! '{ j → j s̲elect₂ M } m̲ap 7 8
75_ ← ("The size of s B A, by rank" "rank" "size" "alpha / r" "alpha / sqrt r" "T itself") ᵖc̲hart! (('{ j → j s̲elect₂ M } m̲ap 5 9) c̲at e̲nclose 8 r̲eshape s̲qrt '+ r̲/ r̲avel T × T)