sourcedemos/lora-lab/lora-lab.xtl
1⍝!/usr/bin/env xetal
2⍝│ LoRA lab: how an adapter's correction is worked out, number by
3⍝│ number, and what its rank changes. A frozen table W of 8 rows and 8
4⍝│ columns, and an adapter of rank r beside it: A (r rows of 8) and B
5⍝│ (8 rows of r). The layer is y = W x + s B (A x), written here as a
6⍝│ pipeline of four stages so that every stage's numbers can be shown:
7⍝│ the input x; the squeeze h = A x (r numbers); the spread u = B h (8
8⍝│ numbers); and the output y = W x + s u. The layer is the one of
9⍝│ X_eTaL-ML's LoRA explorer (its W, A, B and x, made up by fixed
10⍝│ ripples). Nothing is trained here.
11
12"e:" u̲se< "Equations"
13"flow:" u̲se< "Flow"
14"lo:" u̲se< "LoRA"
15
16⍝━━ The layer
17
18⍝│ The rank r (how many numbers the squeeze gives) and alpha; the
19⍝│ strength s is alpha / r.
20r ← 2
21alpha ← 4.0
22s ← alpha ˡᵒs̲caling r
23⍝│ The frozen table W, the adapter's tables A and B, and one input X:
24⍝│ made-up numbers from sines and cosines, so none need be typed.
25c ← f̲loat r̲ange 64
26W ← 8 8 r̲eshape s̲in 0.7 × c × c
27d ← f̲loat r̲ange r × 8
28A ← (r c̲at 8) r̲eshape 0.4 × c̲os 1.1 × d × d
29B ← (8 c̲at r) r̲eshape 0.3 × s̲in 0.9 × d × d
30X ← c̲os 0.4 × 8 t̲ake c
31q ← r − 1
32
33⍝━━ The layer as a pipeline
34
35⍝│ Four stages, a line each: the input, then three equations. A stage
36⍝│ names its value after the $ sign, and a later stage may use it.
37y ← @ ᶠˡᵒʷp̲ipe< c#"
38 X ▹x
39▷ "h[k] | r" ᵉe̲q< "sum j:=0..7 | A[k,j] * x[j]" ▹h
40▷ "u[i] | 8" ᵉe̲q< "sum k:=0..q | B[i,k] * h[k]" ▹u
41▷ "y[i] | 8" ᵉe̲q< "(sum j:=0..7 | W[i,j] * x[j]) + s * u[i]""#
42y
43
44⍝━━ What the adapter changed
45
46⍝│ What the frozen table gives by itself, W x; the correction the
47⍝│ adapter adds to it, s B (A x); and a check that the pipeline's
48⍝│ output is the library's own (lo:f_orward): 1 means yes.
49base ← W '+ '× i̲nner X
50correction ← (ˡᵒd̲elta (A, B, s)) '+ '× i̲nner X
51base
52correction
530.000000000001 > 'm̲ax r̲/ a̲bs y − (W, A, B, s) ˡᵒf̲orward X
54
55⍝━━ What it costs
56
57⍝│ The numbers an adapter of this rank trains on this table, r (8 +
58⍝│ 8), then the 64 of the table itself; and the strength s, which is
59⍝│ alpha divided by the rank.
60(r ˡᵒp̲arams 8 8) c̲at 64
61s
62
63⍝━━ For the page
64
65⍝ The numbers of the layer, row by row, for the pictures.
66r̲avel W
67r̲avel A
68r̲avel B
69⍝ Every product of the three sums, so any number of any stage can be
70⍝ written out: each number of A times the matching number of x (a row
71⍝ for each number of h); each number of B times the matching number of
72⍝ h (a row for each number of u); each number of W times the matching
73⍝ number of x (a row for each number of y).
74r̲avel A × (r c̲at 8) r̲eshape X
75r̲avel B × (8 c̲at r) r̲eshape A '+ '× i̲nner X
76r̲avel W × 8 8 r̲eshape X
ᶠˡᵒʷp̲ipe< expands to
({ @ →
x ← X
h ← "h[k] | r" ᵉe̲q< "sum j:=0..7 | A[k,j] * x[j]"
u ← "u[i] | 8" ᵉe̲q< "sum k:=0..q | B[i,k] * h[k]"
g1:stage ← "y[i] | 8" ᵉe̲q< "(sum j:=0..7 | W[i,j] * x[j]) + s * u[i]"
g1:stage
} @)ᵉe̲q< expands to
(({ @ →
g4:a̲t ← { A i → (1 + i) s̲elect r̲avel A }
g3:c̲heck ← { (A, name, k, from) e →
n ← k s̲elect s̲hape A
bad ← w̲here n̲ot r̲avel (0 ≤ e) ∧ e < n
0 = t̲ally bad ? e
"index" ⎕S̲IGNAL name c̲at " axis " c̲at (f̲ormat k) c̲at ": index " c̲at (f̲ormat (f̲irst bad) s̲elect r̲avel e) c̲at " is outside 0.." c̲at (f̲ormat n − 1) c̲at " (from " c̲at from c̲at ")"
}
g2:r̲ank ← { (A, name, q, ref) →
(t̲ally s̲hape A) = q ? A
"rank" ⎕S̲IGNAL ref c̲at " indexes " c̲at (f̲ormat q) c̲at " axes, but " c̲at name c̲at " has rank " c̲at f̲ormat t̲ally s̲hape A
}
g5:grid ← r
g6:g̲rid ← { k → g5:grid r̲eshape (o̲ffsets k s̲elect g5:grid) 'l̲eft t̲able ('× r̲/ k d̲rop g5:grid) r̲eshape 0 }
g5:grid r̲eshape (({ @ →
g7:grid ← r c̲at ({ v → (0 ≤ v) ? v◆ "bad-equation" ⎕S̲IGNAL "sum's range runs low..high: 0..7" } (1 + 7 − 0))
g8:g̲rid ← { k → g7:grid r̲eshape (o̲ffsets k s̲elect g7:grid) 'l̲eft t̲able ('× r̲/ k d̲rop g7:grid) r̲eshape 0 }
'+ r̲/₂ g7:grid r̲eshape ((
((g2:r̲ank (A, "A", 2, "A[k,j]")) g4:a̲t ((
((A, "A", 2, "j") g3:c̲heck (0 + g8:g̲rid 2)) + (2 s̲elect s̲hape A) ×
((A, "A", 1, "k") g3:c̲heck (0 + g8:g̲rid 1))))) ×
((g2:r̲ank (x, "x", 1, "x[j]")) g4:a̲t (
((x, "x", 1, "j") g3:c̲heck (0 + g8:g̲rid 2))))))
} @))
} @))ᵉe̲q< expands to
(({ @ →
g11:a̲t ← { A i → (1 + i) s̲elect r̲avel A }
g10:c̲heck ← { (A, name, k, from) e →
n ← k s̲elect s̲hape A
bad ← w̲here n̲ot r̲avel (0 ≤ e) ∧ e < n
0 = t̲ally bad ? e
"index" ⎕S̲IGNAL name c̲at " axis " c̲at (f̲ormat k) c̲at ": index " c̲at (f̲ormat (f̲irst bad) s̲elect r̲avel e) c̲at " is outside 0.." c̲at (f̲ormat n − 1) c̲at " (from " c̲at from c̲at ")"
}
g9:r̲ank ← { (A, name, q, ref) →
(t̲ally s̲hape A) = q ? A
"rank" ⎕S̲IGNAL ref c̲at " indexes " c̲at (f̲ormat q) c̲at " axes, but " c̲at name c̲at " has rank " c̲at f̲ormat t̲ally s̲hape A
}
g12:grid ← 8
g13:g̲rid ← { k → g12:grid r̲eshape (o̲ffsets k s̲elect g12:grid) 'l̲eft t̲able ('× r̲/ k d̲rop g12:grid) r̲eshape 0 }
g12:grid r̲eshape (({ @ →
g14:grid ← 8 c̲at ({ v → (0 ≤ v) ? v◆ "bad-equation" ⎕S̲IGNAL "sum's range runs low..high: 0..q" } (1 + q − 0))
g15:g̲rid ← { k → g14:grid r̲eshape (o̲ffsets k s̲elect g14:grid) 'l̲eft t̲able ('× r̲/ k d̲rop g14:grid) r̲eshape 0 }
'+ r̲/₂ g14:grid r̲eshape ((
((g9:r̲ank (B, "B", 2, "B[i,k]")) g11:a̲t ((
((B, "B", 2, "k") g10:c̲heck (0 + g15:g̲rid 2)) + (2 s̲elect s̲hape B) ×
((B, "B", 1, "i") g10:c̲heck (0 + g15:g̲rid 1))))) ×
((g9:r̲ank (h, "h", 1, "h[k]")) g11:a̲t (
((h, "h", 1, "k") g10:c̲heck (0 + g15:g̲rid 2))))))
} @))
} @))ᵉe̲q< expands to
(({ @ →
g18:a̲t ← { A i → (1 + i) s̲elect r̲avel A }
g17:c̲heck ← { (A, name, k, from) e →
n ← k s̲elect s̲hape A
bad ← w̲here n̲ot r̲avel (0 ≤ e) ∧ e < n
0 = t̲ally bad ? e
"index" ⎕S̲IGNAL name c̲at " axis " c̲at (f̲ormat k) c̲at ": index " c̲at (f̲ormat (f̲irst bad) s̲elect r̲avel e) c̲at " is outside 0.." c̲at (f̲ormat n − 1) c̲at " (from " c̲at from c̲at ")"
}
g16:r̲ank ← { (A, name, q, ref) →
(t̲ally s̲hape A) = q ? A
"rank" ⎕S̲IGNAL ref c̲at " indexes " c̲at (f̲ormat q) c̲at " axes, but " c̲at name c̲at " has rank " c̲at f̲ormat t̲ally s̲hape A
}
g19:grid ← 8
g20:g̲rid ← { k → g19:grid r̲eshape (o̲ffsets k s̲elect g19:grid) 'l̲eft t̲able ('× r̲/ k d̲rop g19:grid) r̲eshape 0 }
g19:grid r̲eshape ((({ @ →
g21:grid ← 8 c̲at ({ v → (0 ≤ v) ? v◆ "bad-equation" ⎕S̲IGNAL "sum's range runs low..high: 0..7" } (1 + 7 − 0))
g22:g̲rid ← { k → g21:grid r̲eshape (o̲ffsets k s̲elect g21:grid) 'l̲eft t̲able ('× r̲/ k d̲rop g21:grid) r̲eshape 0 }
'+ r̲/₂ g21:grid r̲eshape ((
((g16:r̲ank (W, "W", 2, "W[i,j]")) g18:a̲t ((
((W, "W", 2, "j") g17:c̲heck (0 + g22:g̲rid 2)) + (2 s̲elect s̲hape W) ×
((W, "W", 1, "i") g17:c̲heck (0 + g22:g̲rid 1))))) ×
((g16:r̲ank (x, "x", 1, "x[j]")) g18:a̲t (
((x, "x", 1, "j") g17:c̲heck (0 + g22:g̲rid 2))))))
} @) + (s ×
((g16:r̲ank (u, "u", 1, "u[i]")) g18:a̲t (
((u, "u", 1, "i") g17:c̲heck (0 + g20:g̲rid 1)))))))
} @))