librarywork/xetal-ml/libs/LoRA/src/LoRA.xtl
LoRA: low-rank adapters: the adapted forward pass in factored form, the merge, the counts, the gradients and steps, LoRA-FA
Import it with an alias of your choice: "lo:" u_se< "LoRA". Put libs/LoRA/src on XETAL_PATH ("just path"); the reference is libs/LoRA/docs. Names with l: are exported; h: names are private to this file.
A frozen weight matrix W (m x n: m outputs, n inputs) is adapted by
a low-rank update s B A: A is r x n, B is m x r, the rank r much
smaller than m and n, s the scale alpha / r. An adapter is the tuple
(W, A, B, s). The inputs are columns: x is one input of n numbers or
an n x N matrix of N inputs, and W x is W '+ '× i̲nner x. The
forward pass keeps the factors apart, B (A x), as training does:
the merged matrix W + s B A is made only to deploy it.
Training moves A and B by the gradients of the squared error (g_rads, s_gd, a_dam through the Optim library); LoRA-FA freezes A (regenerated from a seed, s_eeded) and trains B alone.
The forward pass
ˡs̲caling : Num a => Float -> a -> Float
alpha s̲caling r: the adapter's scale s, alpha / r (so that a
change of rank keeps the update's size).
"lo:" u̲se< "LoRA"
16.0 ˡᵒs̲caling 8 2.0
ˡs̲caling ← { alpha r → alpha ÷ f̲loat r }
ˡs̲calingRs : Num a => Float -> a -> Float
alpha s̲calingRs r: rsLoRA's rank-stabilized scale, alpha / sqrt r
(Kalajdzievski, "A Rank Stabilization Scaling Factor for Fine-Tuning
with LoRA", 2023): with alpha / r the update's learning slows as r
grows; over sqrt r it does not. An adapter takes either, (W, A, B, s).
"lo:" u̲se< "LoRA"
16.0 ˡᵒs̲calingRs 4 8.0
ˡs̲calingRs ← { alpha r → alpha ÷ s̲qrt f̲loat r }
ˡf̲orward : Num a => (a, a, a, a) -> a -> a
(W, A, B, s) f̲orward x: the adapted layer, W x + s B (A x), the
update in its factored form: A x (r numbers per input), then B of
that, never the m x n matrix B A.
"lo:" u̲se< "LoRA"
(2 2 r̲eshape 1.0 0.0 0.0 1.0, 1 2 r̲eshape 1.0 1.0, 2 1 r̲eshape 1.0 -1.0, 0.5) ˡᵒf̲orward 2.0 4.0 5.0 1.0
ˡf̲orward ← { (W, A, B, s) x → (W '+ '× i̲nner x) + s × B '+ '× i̲nner A '+ '× i̲nner x }
The merge
ˡd̲elta : Num a => (a, a, a) -> a
d̲elta (A, B, s): the update itself, s B A, an m x n matrix of
rank r at most.
"lo:" u̲se< "LoRA"
ˡᵒd̲elta (1 2 r̲eshape 1.0 1.0, 2 1 r̲eshape 1.0 -1.0, 0.5) 0.5 0.5 -0.5 -0.5
ˡd̲elta ← { (A, B, s) → s × B '+ '× i̲nner A }
ˡm̲erge : Num a => (a, a, a, a) -> a
m̲erge (W, A, B, s): the adapter folded into the weights, W + s B
A: one matrix again, as fast as the base layer, for one task.
"lo:" u̲se< "LoRA"
ˡᵒm̲erge (2 2 r̲eshape 1.0 0.0 0.0 1.0, 1 2 r̲eshape 1.0 1.0, 2 1 r̲eshape 1.0 -1.0, 0.5) 1.5 0.5 -0.5 0.5
ˡm̲erge ← { (W, A, B, s) → W + ˡd̲elta (A, B, s) }
Gradients
ˡl̲oss : (Float, Float, Float, Float) -> (Float, Float) -> Float
(W, A, B, s) l̲oss (x, y): the squared error of the adapted layer
on the inputs x (columns) against the wanted outputs y, half the sum
of the squared differences, so that its gradient at the output, g,
is the difference itself.
"lo:" u̲se< "LoRA"
(2 2 r̲eshape 1.0 0.0 0.0 1.0, 1 2 r̲eshape 1.0 1.0, 2 1 r̲eshape 1.0 -1.0, 0.5) ˡᵒl̲oss (2 1 r̲eshape 2.0 4.0, 2 1 r̲eshape 4.0 2.0) 1.0
ˡl̲oss ← { a (x, y) → e ← (a ˡf̲orward x) − y 0.5 × '+ r̲/ r̲avel e × e }
ˡg̲rads : Num a => (a, a, a, a) -> (a, a) -> (a, a, a)
(W, A, B, s) g̲rads (x, g): the gradients for the inputs x (an n x
N matrix, one input per column) and g, the loss's gradient at the
output (m x N): with z = A x, dL/dA = s (B' g) x', dL/dB = s g z' and
dL/dx = W' g + s A' (B' g), as the tuple (dA, dB, dx). B' g is made
once and used twice; no m x n matrix but W's own is made.
"lo:" u̲se< "LoRA"
(dA, dB, dx) ← (2 2 r̲eshape 1.0 0.0 0.0 1.0, 1 2 r̲eshape 1.0 1.0, 2 1 r̲eshape 1.0 -1.0, 0.5) ˡᵒg̲rads (2 1 r̲eshape 2.0 4.0, 2 1 r̲eshape 1.0 -1.0)
dB
3.0
-3.0
ˡg̲rads ← { (W, A, B, s) (x, g) → b ← (o̲\ B) '+ '× i̲nner g (s × b '+ '× i̲nner o̲\ x, s × g '+ '× i̲nner o̲\ A '+ '× i̲nner x, ((o̲\ W) '+ '× i̲nner g) + s × (o̲\ A) '+ '× i̲nner b) }
Training steps
ˡs̲gd : Num a => a -> ((a, a, a, a), (a, a)) -> (a, a, a, a)
lr s̲gd ((W, A, B, s), (x, y)): one step of gradient descent on
the squared error (l_oss) at the rate lr, A and B moved by Optim's
s_gd, W and s left as they are: the adapter after the step. p_ower
repeats it.
"lo:" u̲se< "LoRA"
(_, A, B, _) ← 0.1 ˡᵒs̲gd ((2 2 r̲eshape 1.0 0.0 0.0 1.0, 1 2 r̲eshape 1.0 1.0, 2 1 r̲eshape 1.0 -1.0, 0.5), (2 1 r̲eshape 2.0 4.0, 2 1 r̲eshape 4.0 2.0))
B
0.7
-0.7
ˡs̲gd ← { lr ((W, A, B, s), (x, y)) → (dA, dB, _) ← (W, A, B, s) ˡg̲rads (x, ((W, A, B, s) ˡf̲orward x) − y) (W, lr ᵒᵖs̲gd (A, dA), lr ᵒᵖs̲gd (B, dB), s) }
ˡa̲damStart : (Any a, Any b, Num c) => (a, Float, Float, b) -> ((a, Float, Float, b), (Float, Float, Float, Float), c)
a̲damStart (W, A, B, s): what Adam starts from: the adapter, its
four running averages at 0 (A's two, then B's) and the step count 1.
ˡa̲damStart ← { (W, A, B, s) → ((W, A, B, s), (0.0 × A, 0.0 × A, 0.0 × B, 0.0 × B), 1) }
ˡa̲dam : Float -> (((Float, Float, Float, Float), (Float, Float, Float, Float), Float), (Float, Float)) -> ((Float, Float, Float, Float), (Float, Float, Float, Float), Float)
lr a̲dam (((W, A, B, s), (mA, vA, mB, vB), t), (x, y)): Adam's step
t on the squared error, A and B moved by Optim's a_dam at the rate
lr; the next state, the step count one more.
"lo:" u̲se< "LoRA"
((_, _, B, _), _, _) ← 0.1 ˡᵒa̲dam ((ˡᵒa̲damStart (2 2 r̲eshape 1.0 0.0 0.0 1.0, 1 2 r̲eshape 1.0 1.0, 2 1 r̲eshape 1.0 -1.0, 0.5)), (2 1 r̲eshape 2.0 4.0, 2 1 r̲eshape 4.0 2.0))
B
0.9000000003333333
-0.9000000003333333
ˡa̲dam ← { lr (((W, A, B, s), (mA, vA, mB, vB), t), (x, y)) → (dA, dB, _) ← (W, A, B, s) ˡg̲rads (x, ((W, A, B, s) ˡf̲orward x) − y) (A2, mA2, vA2) ← (lr, t) ᵒᵖa̲dam (A, mA, vA, dA) (B2, mB2, vB2) ← (lr, t) ᵒᵖa̲dam (B, mB, vB, dB) ((W, A2, B2, s), (mA2, vA2, mB2, vB2), t + 1) }
LoRA-FA: A frozen
ˡs̲eeded : Num a => a -> Int -> Float
seed s̲eeded r n: an r x n matrix regenerated from a whole-number
seed (LoRA-FA's frozen A, which a checkpoint then need not hold):
sin(seed k k) for k = 1 .. r n, over the square root of n (the
inputs' count, so A x keeps x's size). No chance: the same seed
gives the same matrix everywhere.
"lo:" u̲se< "LoRA"
7 ˡᵒs̲eeded 2 4 0.32849329935939453 0.13545289415393452 0.08367785015140346 -0.44499780218341667 -0.4005672975890204 0.3115061055018265 -0.26829917759428185 0.47413628607807934
ˡs̲eeded ← { seed rn → k ← f̲loat r̲ange '× r̲/ rn rn r̲eshape (s̲in (f̲loat seed) × k × k) ÷ s̲qrt f̲loat 2 s̲elect rn }
ˡg̲radsFA : Num a => (a, a, a, a) -> (a, a) -> (a, a)
(W, A, B, s) g̲radsFA (x, g): LoRA-FA's gradients, (dB, dx): A is
frozen, so its gradient is skipped, but x's still flows back through
A, dx = W' g + s A' (B' g).
ˡg̲radsFA ← { (W, A, B, s) (x, g) → (s × g '+ '× i̲nner o̲\ A '+ '× i̲nner x, ((o̲\ W) '+ '× i̲nner g) + s × (o̲\ A) '+ '× i̲nner (o̲\ B) '+ '× i̲nner g) }
ˡs̲gdFA : Num a => a -> ((a, a, a, a), (a, a)) -> (a, a, a, a)
lr s̲gdFA ((W, A, B, s), (x, y)): LoRA-FA's step of gradient
descent: B alone moves. Every update s B A it reaches has its rows
in A's row space.
ˡs̲gdFA ← { lr ((W, A, B, s), (x, y)) → (dB, _) ← (W, A, B, s) ˡg̲radsFA (x, ((W, A, B, s) ˡf̲orward x) − y) (W, A, lr ᵒᵖs̲gd (B, dB), s) }
ˡa̲damStartFA : (Any a, Any b, Any c, Num d) => (a, b, Float, c) -> ((a, b, Float, c), (Float, Float), d)
a̲damStartFA (W, A, B, s): LoRA-FA's Adam state: the adapter, B's
two running averages at 0 (A has none: half of LoRA's optimizer
memory on a square layer), the step 1.
ˡa̲damStartFA ← { (W, A, B, s) → ((W, A, B, s), (0.0 × B, 0.0 × B), 1) }
ˡa̲damFA : Float -> (((Float, Float, Float, Float), (Float, Float), Float), (Float, Float)) -> ((Float, Float, Float, Float), (Float, Float), Float)
lr a̲damFA (((W, A, B, s), (mB, vB), t), (x, y)): LoRA-FA's Adam
step: B alone moves; the next state.
ˡa̲damFA ← { lr (((W, A, B, s), (mB, vB), t), (x, y)) → (dB, _) ← (W, A, B, s) ˡg̲radsFA (x, ((W, A, B, s) ˡf̲orward x) − y) (B2, mB2, vB2) ← (lr, t) ᵒᵖa̲dam (B, mB, vB, dB) ((W, A, B2, s), (mB2, vB2), t + 1) }
Rank
ʰs̲tep : Num a => Float -> (Float, Int, a) -> (Float, Int, a)
One step of Gram-Schmidt on the rows: the k-th row's remainder R[k], if larger than tol, becomes a direction and is taken out of every row, and the count c goes up.
ʰs̲tep ← { tol (R, k, c) → v ← k s̲elect R n ← s̲qrt '+ r̲/ v × v n ≤ tol ? (R, k + 1, c) (R − (R '+ '× i̲nner v ÷ n) '× t̲able v ÷ n, k + 1, c + 1) }
ˡr̲ank : Num a => Float -> a
r̲ank M: the number of independent rows of the matrix M (its
rank), by Gram-Schmidt: a row counts when what is left of it, after
the directions of the rows before it are taken out, is larger than
1e-9 times the largest row.
"lo:" u̲se< "LoRA"
ˡᵒr̲ank 3 3 r̲eshape 1.0 2.0 3.0 2.0 4.0 6.0 0.0 1.0 1.0 2
ˡr̲ank ← { M → tol ← 0.000000001 × s̲qrt 'm̲ax r̲/ '+ r̲/₂ M × M (_, _, c) ← (f̲irst s̲hape M) '{ R → tol ʰs̲tep R } p̲ower (M, 1, 0) c }
Counts
ˡp̲arams : Num a => a -> a -> a
r p̲arams m n: the numbers an adapter of rank r on an m x n matrix
trains, r (m + n) (B's m r and A's r n), against the m n of the
matrix itself.
"lo:" u̲se< "LoRA"
16 ˡᵒp̲arams 4096 14336 294912
ˡp̲arams ← { r mn → r × '+ r̲/ mn }
ˡe̲xpert : Num a => a -> a -> a
r e̲xpert d f: the numbers rank-r adapters train on one SwiGLU
expert of hidden size d and inner size f: its three matrices (gate
and up, f x d; down, d x f), 3 r (d + f).
"lo:" u̲se< "LoRA"
16 ˡᵒe̲xpert 4096 14336 884736
ˡe̲xpert ← { r df → 3 × r ˡp̲arams df }
ˡe̲xpertFA : Num a => a -> a -> a
r e̲xpertFA d f: the same expert under LoRA-FA: the numbers it
trains (the B's, r (2f + d)), then the frozen A's it regenerates
(r (2d + f)); together r e_xpert d f.
"lo:" u̲se< "LoRA"
16 ˡᵒe̲xpertFA 4096 14336 524288 360448
ˡe̲xpertFA ← { r df → (d, f) ← (f̲irst df, 2 s̲elect df) (r × f + f + d) c̲at r × d + d + f }