Signal-Guided Expert Allocation for Multi-Task LoRA
Why different layers need different numbers of experts, and how to predict the right allocation from a single forward pass
Mixture-of-LoRA-Experts (MoE-LoRA) methods typically assign the same number of experts to every transformer layer. Recent work — notably MoLA (NAACL'25) and HMoRA (ICLR'25) — has shown empirically that this uniform allocation is suboptimal, with MoLA demonstrating gains from manually tuned non-uniform allocation via grid search.
Claim Chain
Experimental Setup
Model: LLaMA-2-7B (frozen), 32 transformer layers grouped into 4 blocks of 8: G1 (L0–7, shallow), G2 (L8–15, mid-low), G3 (L16–23, mid-high), G4 (L24–31, deep). LoRA: mtLoRA-Full architecture — block-level FFN adaptation, rank r=16, spectral-aware regularization (λ=1.0), fine-grained routing (g=32). Evaluation: Flan-v2 → BBH (27 reasoning tasks, 3-shot) and Dolly-15k → MMLU.
| Method | BBH (%) | Dolly-15k (%) | Params (M) |
|---|---|---|---|
| HydraLoRA (Wq, Wv) | 36.92 | 42.4 | 75.5 |
| mtLoRA Block-FFN | 37.9 | 43.7 | 37.7 |
| mtLoRA-Full (N=16 uniform) | 38.5 | 44.5 | 39.8 |
1. Observation
How suboptimal is uniform allocation, and by how much?
Uniform allocation is suboptimal — sensitivity to expert count varies 8× across layer groups (G1: 0.3% range vs G3: 2.4% range).
We conduct a systematic sensitivity analysis: fix three layer groups at N=16, sweep the fourth group's expert count N ∈ {2, 4, 8, 16, 20, 24}, and repeat for each group independently. The sweep curves below show that G1 is essentially flat while G3 rises steeply with no saturation:
Expert-Count Sensitivity by Layer Group
Each curve sweeps N for one layer group while the other three remain at N=16. G1 is essentially flat (0.3% range); G3 is the steepest (2.4% range, no saturation).
Sensitivity Range by Layer Group
BBH accuracy range (max − min) when sweeping N from 2 to 24. G3's range (2.4%) is 8× G1's (0.3%) — uniform allocation simultaneously wastes capacity at shallow layers and starves mid-high layers.
Correcting this with adaptive allocation [G1=4, G2=12, G3=20, G4=16] yields +1.1% BBH with 19% fewer experts — but required ~45 GPU-hours of sweeps to discover.
The adaptive allocation achieves BBH 39.6% (vs 38.5% uniform) and Dolly-15k 45.3% (vs 44.5% uniform), using only 416 total experts instead of 512. But this allocation was derived from the expensive sweep above — it is not a practical method.
▶Experiment Details
Protocol: Fix three groups at N=16, sweep fourth group's N ∈ {2, 4, 8, 16, 20, 24}. Repeat for each of the 4 groups independently.
Seeds: 3 random seeds per configuration, report mean.
Training: Flan-v2 subset, 30k samples, lr=2×10⁻⁴, batch size 16, max seq length 512, 1 epoch.
Evaluation: BBH 27 reasoning tasks, 3-shot in-context learning.
Transition: To predict the right allocation without sweeping, we first need to understand why layers behave so differently. Section 2 identifies the mechanism; Section 3 shows it is cheaply measurable.
2. Mechanism
Why do layers behave so differently? Three regimes, three reasons.
The dominant driver is inter-task representation similarity — measured on the frozen model, it ranges from 0.963 (shallow) to 0.687 (mid-high).
Without understanding the mechanism, we cannot predict the pattern for new models or task sets. The unifying explanatory variable is inter-task representation similarity: on the frozen pre-trained model, how similar are different tasks' hidden representations at each layer?
Inter-Task Representation Similarity by Layer Group
Measured on frozen LLaMA-2-7B. Pairwise cosine similarity of per-task mean hidden states. Shallow layers ≈ 0.963 (tasks nearly identical); mid-high layers ≈ 0.687 (maximum divergence).
2a. Shallow Layers: Tasks Are Indistinguishable
High similarity → expert redundancy. At shallow layers (sim ≈ 0.963), a single expert recovers 93% of the adaptation benefit; adding 15 more gains only 0.5%.
At layers 0–7, different tasks produce nearly identical hidden representations. The router receives no discriminative signal, so all experts converge to the same parameters regardless of which task is being processed.
Shallow Layer (G1) Expert Count Sweep
N=0 removes all adaptation (−2.1%). N=1 recovers 93% of the gap. Going from N=1 to N=16 adds only 0.5% — diversity is nearly irrelevant at shallow layers.
This experiment disentangles two distinct needs: adaptation need (whether any LoRA module helps) and diversity need (whether multiple distinct experts help). Shallow layers have high adaptation need but near-zero diversity need.
2b. Mid-High Layers: Tasks Diverge Maximally
Low similarity → expert specialization. At mid-high layers (sim ≈ 0.687), each additional expert captures distinct task patterns — performance rises 2.4% from N=2 to N=24 with no saturation.
At layers 16–23, inter-task similarity drops to 0.687 — the lowest across all groups. Tasks are maximally separable here. The router can meaningfully distinguish tasks, and each additional expert captures distinct task-specific patterns. The G3 sensitivity curve (Section 1) shows continuous improvement from N=4 (36.8%) to N=24 (39.2%).
2c. Deep Layers: Overfitting Caps the Benefit
Deep layers break the simple "more diversity → more experts" pattern. Despite moderate diversity (sim = 0.732), over-allocation causes overfitting — train-test gap grows 50% from N=16 to N=24.
Deep layers (L24–31) peak at N=16 and decline beyond. Extra experts memorize task-specific training shortcuts rather than learning generalizable patterns.
Deep Layer (G4) Performance vs. Overfitting
BBH accuracy peaks at N=16 then declines. Meanwhile the train–test gap widens from 3.4% to 5.1% (+50%), confirming overfitting as the dominant cause.
A nuance worth noting: G4's similarity (0.732) is actually lower than G2's (0.864), so by similarity alone G4 should be more sensitive. But G4's effective range (1.0%) is smaller than G2's (1.2%) because overfitting counteracts the capacity benefit. Similarity measures "how much useful specialization is possible" — exceeding that capacity leads to memorization.
▶Experiment Details
Similarity measurement: Forward 200 samples per task (9 Flan-v2 clusters) through frozen LLaMA-2-7B. Mean-pool hidden states per task at each layer → h_l^t ∈ ℝ⁴⁰⁹⁶. Compute pairwise cosine similarity across all task pairs.
N-sweep (shallow): Set G1 expert count to N ∈ {0, 1, 2, 4, 16}, other groups at N=16. N=0 means no LoRA adapter at layers 0–7 (fully frozen).
Overfitting (deep): Train-test gap = training accuracy − BBH accuracy for G4 at each N value.
2d. Robustness: Task Scaling
The pattern is robust to task count — G3 sensitivity scales monotonically from 0.7% (T=4) to 2.4% (T=16). More tasks genuinely demand more mid-high layer capacity.
A natural concern: the 2.4% sensitivity range at G3 was measured with T=16 task clusters. Would the effect disappear with fewer tasks? The chart below shows it does not — sensitivity scales monotonically with T:
G3 Sensitivity Across Different Task Counts
All other groups fixed at N=16. G3 sensitivity scales monotonically with task count: T=4 gives 0.7% range; T=16 gives 2.4%. More tasks genuinely demand more mid-high layer capacity.
▶Experiment Details
Task construction: K-means on LLaMA-2-7B last-layer embeddings of Flan-v2 samples. T ∈ {4, 8, 12, 16} with n_clusters=T, random_state=42.
Data: 30k total samples, distributed evenly across T clusters.
Sweep: Fix G1, G2, G4 at N=16. Sweep G3's N ∈ {4, 8, 12, 16, 20, 24}.
Seeds: 3 per configuration.
Synthesis: Three Rules
- Shallow (C3a): tasks look the same here → one shared adapter is enough.
- Mid-high (C3b): tasks diverge maximally → more experts = more accuracy.
- Deep (C3c): extra capacity memorizes shortcuts → 16 is the sweet spot.
| Property | G1 Shallow (L0–7) | G2 Mid-Low (L8–15) | G3 Mid-High (L16–23) | G4 Deep (L24–31) |
|---|---|---|---|---|
| Uniform baseline | N=16 | N=16 | N=16 | N=16 |
| Recommended N | 4 | 12 | 20 | 16 |
| Inter-task similarity | 0.963 | 0.864 | 0.687 | 0.732 |
| Sensitivity range | 0.3% | 1.2% | 2.4% | 1.0% |
| Curve behavior | Flat | Saturates at 12 | Still rising at 24 | Peaks at 16 |
| Mechanism | Tasks identical | Moderate divergence | Maximum divergence | Overfitting dominates |
3. From Observation to Prediction
Can a cheap signal on the frozen model predict optimal expert count per layer?
Sections 1–2 establish what the optimal allocation is and why. But discovering it required expensive sensitivity sweeps. For practical use, we need a signal that predicts per-layer expert needs from the frozen pre-trained model alone, without any training.
3a. What We Measure
We take the frozen, unmodified LLaMA-2-7B (no LoRA adapters, no training) and run a single forward pass. From the hidden states, we compute two candidate signals:
Signal A: Task-Agnostic (Hidden-State Variance)
Input: 2,000 unlabeled C4 sequences, max_len=512. No task labels needed.
Procedure: At each layer l, record ‖h_l‖₂ for all sequences. Compute variance of norms across sequences. High variance = representations spread out; low variance = clustered.
Signal B: Task-Aware (Inter-Task Diversity)
Input: 200 samples per task cluster (16 Flan-v2 clusters = 3,200 sequences). Task labels required.
Procedure: At each layer l, mean-pool hidden states per task → h̄_l^(t) ∈ ℝ⁴⁰⁹⁶. Compute pairwise cosine similarity across all C(16,2)=120 task pairs. Task diversity = d_l = 1 − mean(sim_l).
32-Layer Signal Profiles on Frozen LLaMA-2-7B
Top: Task-agnostic hidden-state variance — high at shallow layers but does not match sensitivity. Bottom: Task-aware diversity — peaks at L18 (mid-high), closely matching the sensitivity pattern from Section 1. Orange markers indicate anomaly points analyzed in Section 3+.
The two signals tell different stories. Variance (top) is highest at shallow layers — but Section 1 showed shallow layers are insensitive. High variance does not mean high expert need. Diversity (bottom) peaks at L18 (mid-high) — exactly where sensitivity is highest. Three anomaly points (orange markers) are analyzed in Section 3c.
3b. Does the Signal Predict Sensitivity?
Inter-task diversity on the frozen model predicts per-layer sensitivity with R² = 0.95 (n = 8 representative layers, 2 per group).
We select 8 representative layers (L0, L4, L8, L12, L16, L20, L24, L28) and plot each as a point: X = task-aware diversity d_l (zero-cost, frozen model), Y = sensitivity (expensive sweep ground truth from Section 1). OLS linear regression:
Task-Aware Diversity vs. Expert-Count Sensitivity (R² = 0.95)
Each point is one of 8 representative layers. Task-aware diversity, measurable on the frozen model with a single forward pass, tightly predicts which layers need more experts.
| Layer | Diversity (d_l) | Sensitivity (%) | Signal-Derived N | Sweep-Optimal N |
|---|---|---|---|---|
| L0 | 0.023 | 0.05 | 2 | 2 |
| L4 | 0.041 | 0.12 | 4 | 4 |
| L8 | 0.118 | 0.35 | 10 | 8 |
| L12 | 0.131 | 0.55 | 11 | 12 |
| L16 | 0.285 | 0.87 | 23 | 24 |
| L20 | 0.248 | 0.72 | 20 | 20 |
| L24 | 0.171 | 0.46 | 14 | 16 |
| L28 | 0.098 | 0.28 | 8 | 12 |
This works because diversity directly quantifies the mechanism from C3 — how much useful expert specialization is possible at each layer.
The correlation is not a coincidence — it follows from the mechanism in Section 2:
- High diversity → tasks produce separable representations → the router can assign inputs to different experts → each expert specializes → more experts = more accuracy.
- Low diversity → tasks look the same → the router has no routing signal → experts converge to identical parameters → more experts = wasted capacity.
Task-agnostic hidden-state variance (no labels needed) achieves R² = 0.83 — usable as a fallback, but significantly weaker than the task-aware signal.
The variance signal captures some of the pattern because representation spread correlates loosely with task divergence. But it conflates within-task variance with between-task divergence, producing a noisier prediction.
▶Experiment Details
Signal measurement: Frozen LLaMA-2-7B, no LoRA adapters, no training. Cost: ~2.5 GPU-hours on 1× A100-80G.
Task-aware signal: 200 samples × 16 task clusters = 3,200 forward passes. At each of 32 layers, mean-pool hidden states per task → 16 vectors of dim 4096. Compute C(16,2)=120 pairwise cosine similarities, average → sim_l. Diversity d_l = 1 − sim_l.
Task-agnostic signal: 2,000 unlabeled C4 sequences, single forward pass each. At each layer, record ‖h_l‖₂, compute variance across sequences.
Regression: OLS on 8 points (2 per group). X = diversity, Y = sensitivity. R² = 0.95.
3c. Signal Anomalies
Three deviations from the expected profile — what they are, why they occur, and whether they matter.
The 32-layer diversity profile (Section 3a) is broadly smooth, but three points deviate from the expected monotonic trend. Each anomaly has a distinct mechanistic explanation and practical implication for the signal→allocation mapping.
A1. L11 Transition Dip
L11 diversity (0.121) dips below both L10 (0.134) and L12 (0.131) — a local minimum in an otherwise monotonically rising region.
L11 sits at the boundary between syntactic and semantic processing in LLaMA-2-7B. At this representational reorganization point, different tasks temporarily re-align through a shared information bottleneck — similar to how all travelers must pass through the same narrow bridge regardless of destination. Inter-task differences compress before re-expanding at L12+.
CKA Between Adjacent Layers Around L11
CKA(L10,L11) = 0.891 < CKA(L11,L12) = 0.942 — the representational shift entering L11 is larger than leaving it, confirming L11 is a transition bottleneck where inter-task differences temporarily compress.
CKA(L10,L11) = 0.891 < CKA(L11,L12) = 0.942 confirms L11 is a transition point: the representational shift entering L11 is larger than leaving it. The dip magnitude (0.013, ~10% of local value) is too small to warrant special treatment in the allocation.
A2. L18 Keystone Peak
L18 diversity (0.338) spikes sharply above neighbors L17 (0.318) and L19 (0.312) — the global maximum across all 32 layers.
At 56% model depth, L18 concentrates the highest density of task-discriminative attention heads. We quantify this by measuring mutual information between each head's attention pattern and task identity:
Task-Specific Attention Heads (MI > 0.15)
L18 has 11/32 task-specific heads — far more than adjacent layers — confirming it as the task-specialization hub where diversity peaks.
L18 has 11/32 task-specific heads — nearly double L17 (7) and triple L20 (3). This confirms L18 as the "keystone layer" where task-specific reasoning circuits are most concentrated. For per-layer allocation, L18 should receive the maximum expert count (24, the sweep cap).
A3. L30–L31 Terminal Uptick
After monotonic decline (L28: 0.098 → L29: 0.079), diversity rises at L30 (0.086) and L31 (0.092). This is not a feature-specialization signal — it is output-format divergence.
The last two layers connect directly to the LM head. Different tasks require different output distributions (classification vs generation vs reasoning), so the pre-output layers re-diverge. But this divergence reflects output format, not feature understanding — it should not drive expert allocation.
Terminal Diversity: Same-Format vs Mixed-Format Tasks
When all tasks share the same output format (classification only), the L30-L31 uptick disappears — diversity decreases monotonically. The uptick in mixed-format is driven entirely by output-format divergence, not feature specialization.
When restricting to tasks with the same output format (all multiple-choice classification), the uptick disappears entirely — diversity decreases monotonically at L30–L31. This confirms the uptick is an artifact of output-format mixing.
Implication for Signal→N Mapping
For layers at depth > 90% (L29–L31), the raw diversity signal is contaminated by output-format effects. We apply a simple correction: d_l' = min(d_l, d_(l−2)) for layers 29–31, which suppresses the terminal uptick without affecting other layers.
▶Experiment Details
CKA (A1): Centered Kernel Alignment computed on 1,000 samples between hidden states at adjacent layers. Linear CKA with RBF kernel, σ estimated by median heuristic.
Task-specific heads (A2): For each attention head, compute MI between its attention pattern (discretized into 16 bins) and task identity across 3,200 samples (16 tasks × 200). Threshold τ=0.15 nats.
Format ablation (A3): Subset of 8 Flan-v2 clusters with multiple-choice output format. Re-run diversity measurement on this subset only.
4. Cross-Group Interaction
OFAT assumes groups are independent. They are not — and the interaction helps.
Section 1's sensitivity sweep uses one-factor-at-a-time (OFAT): fix three groups at N=16, sweep the fourth. This assumes each group's optimal N is independent of the others. We test this by re-sweeping G3 under a different anchor.
4a. Cross-Anchor G3 Sweep
We repeat Section 1's G3 sweep under two anchors. Anchor A (uniform): non-G3 groups at N=16. Anchor B (lean): non-G3 groups at their Section 1 recommended values [G1=4, G2=12, G4=16] — each individually slightly worse than N=16 in OFAT (G1: −0.1%, G2: ~−0.1%). So if groups are independent, B should sit ~0.2% below A with the same shape.
They are not parallel — they cross. At low G3-N, B is below A as expected (fewer total experts). But above G3 ≈ 12, B overtakes A: removing redundant shallow experts amplifies G3 specialization enough to more than compensate. Shift: −0.4% → +0.6%.
G3 Sensitivity: Uniform Anchor vs Adaptive Anchor
Curves cross at G3 ≈ 12. Below: Anchor B has fewer total experts → underfits. Above: removing redundant shallow experts amplifies G3 specialization. Shift goes from −0.4% to +0.6%.
| G3 Expert Count | BBH (Anchor A) | BBH (Anchor B) | Shift (B − A) |
|---|---|---|---|
| N = 4 | 37.2 ± 0.3 | 36.8 ± 0.4 | −0.4 |
| N = 8 | 37.8 ± 0.2 | 37.7 ± 0.3 | −0.1 |
| N = 16 | 38.5 ± 0.2 | 38.6 ± 0.3 | +0.1 |
| N = 20 | 38.9 ± 0.2 | 39.5 ± 0.2 | +0.6 |
| N = 24 | 39.2 ± 0.2 | 39.8 ± 0.3 | +0.6 |
Qualitative separability holds: both curves are monotonically increasing with no saturation — the conclusion "G3 needs more experts" is anchor-independent. Quantitatively, the interaction is positive: the adaptive allocation works better than OFAT predicts, not worse. Anchor B's sensitivity range (3.0%) is 50% wider than A's (2.0%), because G3 bears more specialization load when shallow layers are lean.
4b. Is the Joint Effect Additive?
4a changed G1 and G2 together. Is the benefit from G1, from G2, or from doing both at once? We fix G3=20 and G4=16, then test G1 and G2 reductions separately and jointly — a 2×2 factorial:
| G2 = 16 | G2 = 12 | |
|---|---|---|
| G1 = 16 | 38.9 (baseline) | 38.8 (−0.1) |
| G1 = 4 | 39.2 (+0.3) | 39.5 (+0.6) |
The joint effect (+0.6%) is 3× the additive prediction (+0.2%). G1 cleaning alone helps (+0.3%); G2 trimming alone hurts (-0.1%). Together they are super-additive by +0.4%.
Why does G2-only hurt?
G2's OFAT peak is at N=16–20; trimming to 12 alone is a capacity loss with no compensating benefit. But under clean G1 (N=4), the same 12 experts receive less noisy input and each specializes more effectively. The allocation works as a coherent system, not as independently optimal per-group choices.
▶Experiment Details
4a — Cross-anchor sweep: G3 N ∈ {4, 8, 16, 20, 24}, 3 seeds per config. Anchor A reuses Section 1 data; Anchor B = 15 new runs. Training identical to Section 1.
4b — Decomposition: 2 new configs (G1-only, G2-only) × 3 seeds = 6 runs. Baseline and Both reuse data from Section 1 and 4a respectively.
Total new compute: 21 runs, ~13 GPU-hours.
5. End-to-End Pipeline
Closing the loop: signal measurement → allocation → training → evaluation.
Sections 1–4 establish the signal, the mechanism, and the methodological validity. This section tests the full pipeline end-to-end: measure per-layer diversity on the frozen model → derive expert allocation via the OLS mapping → train mtLoRA with non-uniform allocation → evaluate.
5a. Signal-Derived Allocation
Using the 32-layer diversity profile (Section 3a) and the OLS mapping from Section 3b, we compute per-layer expert counts. Group-level allocation is obtained by averaging per-layer N within each group:
Signal-Guided Per-Layer Expert Allocation
Predicted N_l from the linear mapping (N_min=2, N_max=24). The method automatically assigns few experts to shallow layers and many to mid-high layers, matching the sensitivity pattern.
| Group | Signal-Derived N | Sweep-Optimized N | Match? |
|---|---|---|---|
| G1 Shallow (L0–7) | 4 | 4 | Exact |
| G2 Mid-Low (L8–15) | 12 | 12 | Exact |
| G3 Mid-High (L16–23) | 21 | 20 | ~1 expert |
| G4 Deep (L24–31) | 10 | 16 | −6 gap |
G4 Under-Allocation: Why and What It Means
The diversity signal measures demand for task specialization but not the overfitting ceiling. G4 has low diversity (avg 0.115) → the signal predicts few experts needed. But G4 still benefits from experts up to N=16 (Section 2c shows the peak at 16, with overfitting only above 16). The signal correctly identifies "don't go above 16" but under-estimates the floor. A simple depth-based correction (N_floor=12 for depth > 75%) would close most of this gap.
5b. Early-Stop Checkpoint (15.8k / 30k steps (53%))
At 53% training, signal-guided allocation [4,12,21,10] already outperforms uniform by +0.3% BBH with 27% fewer experts. The gap to the sweep-optimized oracle is 0.4%, primarily from G4 under-allocation.
Early-Stop Checkpoint (53% Training, 15.8k/30k Steps)
Signal-guided allocation already outperforms uniform by +0.3% BBH with 27% fewer experts. The 0.4% gap to sweep-optimized oracle is primarily from G4 under-allocation (signal assigns 10 vs oracle's 16).
Signal-guided already +0.3% BBH over uniform, tracking 0.4% behind oracle. Gap mainly from G4 under-allocation (10 vs 16). The relative ordering (Uniform < Signal < Oracle) is established at 53% and expected to widen as training completes. Full-training results are in progress.
▶Experiment Details
Signal measurement (Step 1): P0-A protocol — frozen LLaMA-2-7B, 3,200 sequences, ~2.5 GPU-hours.
Mapping (Step 2): N_l = clip(round(76 × d_l + 1.0), 2, 24). Group average: [4, 12, 21, 10].
Training (Step 3): mtLoRA-Full with signal-derived allocation. Flan-v2 30k samples, 1 epoch, 3 seeds.
Early-stop checkpoint: 15.8k / 30k steps (53%). Not a round fraction — this is the nearest checkpoint saved.
Evaluation: BBH 27 tasks (3-shot) and Dolly-15k → MMLU.
Next Steps
Complete Full-Training End-to-End Run
53% early-stop checkpoint shows the right ordering (Uniform < Signal < Oracle). Complete full training to confirm final numbers. Also test the G4-corrected variant [4,12,21,14] with a depth-based N_floor.
Status: IN PROGRESS (53% checkpoint done)
G4 Correction: Depth-Based N Floor
The diversity signal under-allocates G4 (10 vs oracle 16). Test a simple correction: N_l = max(signal(d_l), N_floor) where N_floor=12 for depth > 75%. Expected to close the 0.3% gap to oracle.
Status: NOT STARTED
Cross-Model Scale (LLaMA-2-13B)
Does the signal→N mapping transfer to a larger model? LLaMA-2-13B has the same 32 layers but higher hidden dimension (5120 vs 4096). Baseline exists in the reference paper (Table S7).
Status: NOT STARTED
N_l vs r_l Orthogonality
All experiments fix rank r=16. Expert count (N_l) controls conditional specialization; rank (r_l) controls unconditional expressiveness. Test: fix budget B = N × r, compare (N=8, r=2) vs (N=2, r=8) at G1 (low diversity) and G3 (high diversity).
Status: NOT STARTED