Simo76 commited on
Commit
3f5b03f
·
1 Parent(s): 64f78fc

Add architecture documentation for Nested LoRA

Browse files

Document the architecture and functionality of Nested LoRA, including the problem of cold starts on rank transitions, the solution using a single adapter pair, and the controller's trajectory-based approach.

Files changed (1) hide show
  1. docs/architecture.md +120 -0
docs/architecture.md ADDED
@@ -0,0 +1,120 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Architecture — Nested Orbital LoRA
2
+
3
+ ## The problem: cold start on rank transitions
4
+
5
+ Standard multi-rank LoRA keeps separate adapter matrices per rank level:
6
+
7
+ ```
8
+ r=4: A4(4, d), B4(d, 4)
9
+ r=8: A8(8, d), B8(d, 8)
10
+ r=16: A16(16,d), B16(d, 16)
11
+ ```
12
+
13
+ When the controller switches from r=16 to r=8, the r=8 adapter has independent weights that never benefited from training at r=16. Each transition is a partial cold start. This caused 3-6 point F1 loss vs baseline in our experiments (V1-V4).
14
+
15
+ ## The solution: one particle, multiple orbitals
16
+
17
+ Nested LoRA uses a single adapter pair with the maximum rank. Smaller ranks are obtained by slicing:
18
+
19
+ ```
20
+ A(16, d) and B(d, 16) ← one pair, always present
21
+
22
+ r=4: A[:4, :], B[:, :4] ← first 4 dimensions
23
+ r=8: A[:8, :], B[:, :8] ← first 8 dimensions
24
+ r=16: A[:16,:], B[:, :16] ← full matrix
25
+ ```
26
+
27
+ The metaphor: one particle that can occupy different energy orbitals. When descending from r=16 to r=4, dimensions 0-3 retain everything they learned. Dimensions 4-15 are paused (no gradient), not destroyed. When ascending back, they resume exactly where they left off.
28
+
29
+ ```
30
+ r4 ⊂ r8 ⊂ r16
31
+ ```
32
+
33
+ ### Scaling
34
+
35
+ To maintain consistent output magnitude across ranks, the output is scaled by `max_rank / active_rank`:
36
+
37
+ ```python
38
+ scale = 16 / r
39
+ output = base + delta * scale
40
+ ```
41
+
42
+ At r=4 the scale is 4.0 (amplify the smaller subspace). At r=16 the scale is 1.0 (no amplification). This is analogous to the alpha/r scaling in standard LoRA.
43
+
44
+ ## Controller: trajectory with orbital memory
45
+
46
+ ### From threshold controller to trajectory controller
47
+
48
+ Early versions used threshold-based FSM: if φ > θ₁, switch to r=16. This had two problems: static thresholds don't generalize across tasks/models, and the controller oscillated or got stuck.
49
+
50
+ The orbital controller replaces thresholds with a trajectory:
51
+
52
+ ```
53
+ Ascend: stress detected → jump to higher orbital, push delta to stack
54
+ Hold: oscillating → stay, don't move
55
+ Descend: confirmed stable → pop delta, symmetric return
56
+ ```
57
+
58
+ The orbit stack records the exact sequence of jumps. When descending, the controller reverses them in order, ensuring symmetric return to the previous state.
59
+
60
+ ### Stress signal
61
+
62
+ ```
63
+ φ(t) = |loss - EMA(loss)| + 2.0 × max(0, loss - prev_loss)
64
+ ```
65
+
66
+ Two components:
67
+ - **Deviation from trend**: catches sustained instability
68
+ - **Spike detection**: catches sudden deterioration (weighted 2x)
69
+
70
+ ### Adaptive thresholds
71
+
72
+ ```
73
+ t_stress = μ(φ_recent) + 0.7σ
74
+ t_stable = max(μ(φ_recent) - 0.3σ, 0)
75
+ ```
76
+
77
+ Auto-calibrate to loss scale. No manual tuning needed.
78
+
79
+ ### Stability confirmation
80
+
81
+ Descent requires `stable_window` consecutive steps below `t_stable`. This prevents premature return after a brief lull during ongoing instability.
82
+
83
+ ## Lifecycle of a training run
84
+
85
+ ```
86
+ Step 0 Initialize at max rank (warmup)
87
+ Step 1-W Build EMA baseline, accumulate φ history
88
+ Step W+1 Drop to ground state (r=4)
89
+ Step W+2... Controller active:
90
+ stress → ascend (push delta)
91
+ stable → descend (pop delta)
92
+ neutral → hold
93
+ ```
94
+
95
+ ### Warmup rationale
96
+
97
+ The controller needs calibrated thresholds before making decisions. Training at max rank during warmup ensures the model makes initial progress regardless of controller behavior.
98
+
99
+ ### Ground state rationale
100
+
101
+ After warmup, dropping to the minimum rank tests whether the task actually needs high capacity. If it does, the stress signal will push the controller back up immediately. If it doesn't, the system saves rank.
102
+
103
+ ## Comparison with existing methods
104
+
105
+ | Property | Standard LoRA | AdaLoRA | Unified-LoRA |
106
+ |----------------------|---------------|--------------------|-----------------------|
107
+ | Rank control | Fixed | SVD importance | Stress feedback |
108
+ | Control type | None | Open-loop | Closed-loop |
109
+ | Shock reaction | None | Indirect | Immediate |
110
+ | Transition cost | N/A | SVD per step | O(1) slice |
111
+ | Architecture | Single rank | Pruned rank | Nested orbitals |
112
+ | Black-box compatible | Yes | No (needs grads) | Yes |
113
+ | Overhead per step | 0 | O(r² × layers) | O(1) |
114
+
115
+ ## Limitations
116
+
117
+ - Validated on DistilBERT (67M). Scale to 7B+ not yet confirmed.
118
+ - The 15% rank saving on DistilBERT is small in absolute compute terms. The value proposition strengthens at larger scale where rank savings translate to meaningful memory/time reduction.
119
+ - On perfectly stable training, the controller adds no value (but causes no harm).
120
+ - The orbit stack can grow unboundedly in theory, though in practice it stays shallow (1-3 entries).