Repository navigation
fix(sona): make MicroLoRA learn from feedback and make EWC++ take effect - #1258
Merged
Merged
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
SONA's MicroLoRA and EWC++ do not learn as described, so
applyLoracannot be used as an adaptive layer, and EWC++ never protects anything.MicroLoRA (
crates/sona/src/lora.rs)accumulate_gradientaddedgradient_estimate * qualityto every rank ofup_projwithout using the input, and never traineddown_proj. The update was the same for every input. Its effect onforward(x)took the sign of the fixeddown · xterm 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.|up_proj|reached 7 after 50 steps and 707 after 5,000 at lr 0.1.down_projstarted at ±0.01. Updates scale with(down · x)², so learning stalled.recordStepuses 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)LoopCoordinatorsizedEwcPlusPlusashidden_dim * base_lora_rank * 2, but the background loop passes gradients withembedding_dimentries.update_fisher,detect_task_boundaryandapply_constraintsall return early on the length mismatch. After 120 forced cycles with a distribution shift,samples_seenandewc_taskswere still 0.E[g²].apply_constraintsonly shrinks gradients on important weights. Nothing pulled weights back toward an earlier task's optimum, andset_optimal_weightswas never called.apply_micro_lorahad no EWC++ at all.Change
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 initialisedown_projwith the usual1/sqrt(hidden_dim)bound. New helpers:take_accumulated,apply_gradient,flat_weights,set_flat_weights.embedding_dim. Store a bias-corrected Fisher when a task ends; behaviour within a task is unchanged. NewEwcPlusPlus::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 = 0turns it off.start_new_task()onInstantLoop,LoopCoordinatorandSonaEngine.CoordinatorStats.instant_ewc_tasks(serde default 0). WASM:recordStepWithActivations(trajectoryId, nodeId, activations, score)andstartNewTask();getStats()includesinstant_ewc_tasks.runInstantCycle()after feedback), and that task boundaries should be marked explicitly.Behaviour changes
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. Setewc_lambdalower to favour the new task.micro_lora_lrof 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):startNewTask()startNewTask()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:
main(46d816d):applyLoramoves (|applyLora(x) − x|up to 0.13|x|), butrecordStepcan 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_lambdachanges nothing, andewc_tasksstays 0.@ruvector/edge-full0.1.0 SONA build: theapplyLoradelta 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 ofmainit stops at the first check, becauserecordStepWithActivationsdoes not exist there yet.cargo fmt --checkpasses.cargo clippy -p ruvector-sona --all-targets --no-deps -- -D warningspasses, natively and forwasm32-unknown-unknownwith--features wasm.wasm-pack build --target web --release --features wasmbuilds.To reproduce:
cd crates/sona wasm-pack build --target web --release --features wasm node tests-js/lora-ewc.mjsNot changed
startNewTask().recordStepWithActivationsorstartNewTaskyet.MicroLoRA/BaseLoRA/LoraConfigclasses in the prebuiltexamples/edge-full/pkg/rvliteare a different implementation. I could not find their Rust source in the repository, so they are unchanged.trajectory::tests::test_builderasserts a latency above 0 µs and failed once on a busy machine; it passed on rerun. It is unrelated to this change.