Skip to content

fix(sona): make MicroLoRA learn from feedback and make EWC++ take effect - #1258

Merged
ruvnet merged 7 commits into
ruvnet:mainfrom
shaal:fix/microlora-ewc
Oct 24, 2026
Merged

ruvnet merged 7 commits into
ruvnet:mainfrom
shaal:fix/microlora-ewc

Conversation

@shaal

@shaal shaal commented Oct 24, 2026 •

Copy link
Copy Markdown
Contributor

Problem

SONA's MicroLoRA and EWC++ do not learn as described, so applyLora cannot be used as an adaptive layer, and EWC++ never protects anything.

MicroLoRA (crates/sona/src/lora.rs)

  • accumulate_gradient added gradient_estimate * quality to every rank of up_proj without using the input, and never trained down_proj. The update was the same for every input. Its effect on forward(x) took the sign of the fixed down · x term instead of the feedback: in a check with 40 random inputs, only 17 single updates moved the output toward the feedback direction. Two inputs trained toward different targets got exactly opposite deltas.
  • Nothing bounded the weights. |up_proj| reached 7 after 50 steps and 707 after 5,000 at lr 0.1.
  • down_proj started at ±0.01. Updates scale with (down · x)², so learning stalled.
  • In WASM, recordStep uses the query embedding as the step activation, so from JS the learning signal could only point back along the query.

EWC++ (crates/sona/src/ewc.rs, loops/coordinator.rs)

  • LoopCoordinator sized EwcPlusPlus as hidden_dim * base_lora_rank * 2, but the background loop passes gradients with embedding_dim entries. update_fisher, detect_task_boundary and apply_constraints all return early on the length mismatch. After 120 forced cycles with a distribution shift, samples_seen and ewc_tasks were still 0.
  • The online Fisher is an EMA that starts at zero, with no bias correction. After 50 samples the estimate was about 5% of E[g²].
  • apply_constraints only shrinks gradients on important weights. Nothing pulled weights back toward an earlier task's optimum, and set_optimal_weights was never called.
  • The MicroLoRA used by apply_micro_lora had no EWC++ at all.

Change

  • MicroLoRA: back-propagate the gradient through delta = scale · upᵀ(down · x), so both factors train and the update depends on the input. Clip each update (MICRO_LORA_MAX_GRAD_NORM), cap the adapter gain (MICRO_LORA_MAX_GAIN) so |applyLora(x) − x| ≤ |x|, and initialise down_proj with the usual 1/sqrt(hidden_dim) bound. New helpers: take_accumulated, apply_gradient, flat_weights, set_flat_weights.
  • EWC++: size the coordinator's EWC++ to embedding_dim. Store a bias-corrected Fisher when a task ends; behaviour within a task is unchanged. New EwcPlusPlus::apply_anchor, the closed-form (proximal) step for the EWC penalty, which cannot overshoot for any lambda or learning rate. The instant loop gets its own EWC++ over the flat MicroLoRA weights: each flush constrains and anchors the update, updates the Fisher, and records optimal weights at task boundaries. ewc_lambda = 0 turns it off.
  • API: start_new_task() on InstantLoop, LoopCoordinator and SonaEngine. CoordinatorStats.instant_ewc_tasks (serde default 0). WASM: recordStepWithActivations(trajectoryId, nodeId, activations, score) and startNewTask(); getStats() includes instant_ewc_tasks.
  • README: documents the new calls, the output bound, a learning rate that works for unit-length embeddings, that updates apply when the instant loop flushes (call runInstantCycle() after feedback), and that task boundaries should be marked explicitly.

Behaviour changes

  • With ewc_lambda > 0 (the built-in configs use 1000 or 2000), EWC++ now takes effect. After a task boundary, weights that mattered for the previous task adapt more slowly. Set ewc_lambda lower to favour the new task.
  • The MicroLoRA delta is bounded by the input norm, and single updates are clipped.
  • With the default micro_lora_lr of 0.002, learning from unit-length embeddings is slow. About 0.02 learns within a few hundred feedback calls. The default is unchanged.

Results

tests-js/lora-ewc.mjs, 64 dims: task A's held-out alignment rises from 0.12 to 0.72. After task B trains, it is 0.15 without EWC++ (task B reaches 0.65) and 0.64 with lambda 2000 (task B reaches 0.06).

A separate check (32 dims, different data and seeds, lr 0.05, runInstantCycle() after each feedback, task B's inputs at cosine 0.99 to task A's, different target):

task A after A task A after B task B
no EWC++ 1.00 −0.14 1.00
lambda 100 + startNewTask() 1.00 0.52 0.61
lambda 2000 + startNewTask() 1.00 0.92 −0.03

With task B's inputs at cosine 0.77 to task A's, the rows are nearly the same. After 400 feedback calls at lr 0.5, |applyLora(x) − x| / |x| stays at 0.20 or below.

Before this change, the same check from JS:

  • A build of main (46d816d): applyLora moves (|applyLora(x) − x| up to 0.13 |x|), but recordStep can only use the query as the step activation, so feedback cannot aim it at a target. Alignment with the targets stays at 0.02, ewc_lambda changes nothing, and ewc_tasks stays 0.
  • The published @ruvector/edge-full 0.1.0 SONA build: the applyLora delta stays exactly zero after 400 feedback calls.

Tests

  • cargo test -p ruvector-sona: 114 passed, 10 of them new. They cover updates that move toward the feedback and depend on the input, the gain bound, the flat-weight round trip, background cycles feeding EWC++, the bias-corrected Fisher, apply_anchor, task boundaries in both loops, and reduced forgetting.
  • crates/sona/tests-js/lora-ewc.mjs: an end-to-end check against the built web package, 6 of 6 pass. Against a build of main it stops at the first check, because recordStepWithActivations does not exist there yet.
  • cargo fmt --check passes. cargo clippy -p ruvector-sona --all-targets --no-deps -- -D warnings passes, natively and for wasm32-unknown-unknown with --features wasm.
  • wasm-pack build --target web --release --features wasm builds.

To reproduce:

cd crates/sona
wasm-pack build --target web --release --features wasm
node tests-js/lora-ewc.mjs

Not changed

  • Automatic task-boundary detection is conservative. In the check above, switching from task A to task B registered no task, so the README recommends startNewTask().
  • The NAPI bindings do not expose recordStepWithActivations or startNewTask yet.
  • The MicroLoRA / BaseLoRA / LoraConfig classes in the prebuilt examples/edge-full/pkg/rvlite are a different implementation. I could not find their Rust source in the repository, so they are unchanged.
  • trajectory::tests::test_builder asserts a latency above 0 µs and failed once on a busy machine; it passed on rerun. It is unrelated to this change.

shaal and others added 7 commits October 23, 2026 23:12
accumulate_gradient added gradient_estimate * quality to every rank of
up_proj without using the input, and never trained down_proj. The update
was the same for every input, and its effect on forward(x) took the sign
of the fixed (down . x) term instead of the feedback. In a check with 40
random inputs, only 17 single updates moved the output toward the
feedback direction. Two inputs trained toward different targets got
exactly opposite deltas. Repeated feedback also grew up_proj linearly
(|up| 7 after 50 steps and 707 after 5,000 at lr 0.1).

- Back-propagate gradient_estimate through delta = scale * up^T (down x),
  so both factors are trained and the update depends on the input.
- Clip each update to MICRO_LORA_MAX_GRAD_NORM and cap the adapter gain
  at MICRO_LORA_MAX_GAIN, so |forward(x)| <= |x|.
- Initialise down_proj with the usual 1/sqrt(hidden_dim) bound instead
  of 0.01. Updates scale with (down . x)^2, so the old bound stalled
  learning.
- Add take_accumulated, apply_gradient, flat_weights and
  set_flat_weights, so callers can work on the flat gradient.

With the change all 40 updates move toward the feedback, each input
moves toward its own target, and the gain stays at or below 1.
EWC++ had no effect in practice:

- LoopCoordinator sized EwcPlusPlus as hidden_dim * base_lora_rank * 2,
  but the background loop passes gradients with embedding_dim entries.
  update_fisher, detect_task_boundary and apply_constraints all return
  early on the length mismatch. After 120 forced cycles with a
  distribution shift, samples_seen and ewc_tasks were still 0.
- The online Fisher is an EMA that starts at zero. Without bias
  correction, the estimate after 50 samples was about 5% of E[g^2]
  (0.012 for a true 0.25).
- apply_constraints only shrinks gradients on important weights. Nothing
  pulled weights back toward an earlier task's optimum, and
  set_optimal_weights was never called.
- The MicroLoRA used by apply_micro_lora had no EWC++ at all.

Changes:

- Size the coordinator's EWC++ to embedding_dim.
- Store a bias-corrected Fisher when a task ends. Behaviour within a
  task is unchanged.
- Add EwcPlusPlus::apply_anchor, the closed-form (proximal) step for the
  EWC penalty. It cannot overshoot for any lambda or learning rate.
- Give the instant loop its own EWC++ over the flat MicroLoRA weights.
  Each flush constrains and anchors the update, updates the Fisher and
  checks for a task boundary; optimal weights are recorded at each
  boundary. ewc_lambda = 0 turns it off.
- Add start_new_task() to InstantLoop, LoopCoordinator and SonaEngine,
  and instant_ewc_tasks to CoordinatorStats (serde default 0).

Check with a rank-2 MicroLoRA, where task B's inputs overlap task A's
and have a different target: without EWC++, task A's held-out alignment
fell from 0.67 to 0.00 while task B trained. With lambda 2000 it fell
from 0.66 to 0.56, and task B reached 0.21 instead of 0.72. Lower
lambda trades less protection for faster learning on the new task.
recordStep uses the query embedding as the step activation. From JS the
learning signal could only point back along the query, so the MicroLoRA
could not learn to move one input toward another.

- Add recordStepWithActivations(trajectoryId, nodeId, activations,
  score).
- Add startNewTask() to mark an EWC++ task boundary.
- getStats() now includes instant_ewc_tasks.
- Add tests-js/lora-ewc.mjs, an acceptance test against the web build
  (6 checks), and document the API in the README.

Acceptance run: task A's held-out alignment rises from 0.12 to 0.72.
After task B trains, it is 0.15 without EWC++ and 0.64 with lambda
2000.
Automatic EWC++ task detection is conservative. In a check where only
the feedback targets changed between two tasks, it registered no task,
so the README now recommends startNewTask() instead of implying the
boundary is always found.
From JS, feedback without runInstantCycle() changes applyLora output only
occasionally. The README example already calls it; say why.
…al MicroLoRA/EWC++ change

The behavioral parity harness still passes. The rust-sona fingerprint moved
(L2 7.2e-7 against a 2.5e-8 tolerance; unmodified main measures 5.5e-9) because
this PR deliberately changes MicroLoRA learning and EWC++ sizing. Regenerated
with scripts/sona-drift/rvf-fingerprint.mjs --update after building
npm/packages/ruvector so the ruvector-cli entry is preserved; ids and sidecar
are unchanged.

Co-Authored-By: claude-flow <ruv@ruv.net>
@ruvnet
ruvnet merged commit f148455 into ruvnet:main Oct 24, 2026
59 of 61 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants