diff -r df92e78cc3f4 -r 677a5fd1b014 src/fb.rs
--- a/src/fb.rs	Thu Jul 16 19:51:51 2026 +0300
+++ b/src/fb.rs	Sun Jul 19 07:30:06 2026 +0200
@@ -188,17 +188,14 @@
 
     // Run the algorithm
     for state in iterator.iter_init(|| full_stats(&μ, ε, stats.clone())) {
+        let maybe_μ_base = config.merge_now(&state).then(|| μ.clone());
+        let μ_base_len = μ.len();
+
         // Calculate smooth part of surrogate model.
-        // TODO: optimise τ to be applied to residual.
         let mut τv = f.differential(&μ) * τ;
 
-        // Save current base point for merge
-        let μ_base_len = μ.len();
-        let maybe_μ_base = config.merge_now(&state).then(|| μ.clone());
-
-        // Insert and reweigh
-        let (maybe_d, _within_tolerances) = prox_penalty
-            .insert_and_reweigh(&mut μ, &mut τv, τ, ε, config, &reg, &state, &mut stats)?;
+        // Do spike insertiona nd finite-dimensional weight optimisation
+        prox_penalty.insert_and_reweigh(&mut μ, &mut τv, τ, ε, config, &reg, &state, &mut stats)?;
 
         stats.inserted += μ.len() - μ_base_len;
 
@@ -218,12 +215,22 @@
 
         stats.pruned += prune_with_stats(&mut μ);
 
+        // Do extra weight optimisation step heuristic
+        for _ in 0..config.extra_weight_optimisation_steps {
+            τv = f.differential(&μ) * τ;
+            prox_penalty.reweigh(&mut μ, &mut τv, τ, ε, config, &reg, &state, &mut stats)?;
+        }
+
+        if config.extra_weight_optimisation_steps > 0 {
+            stats.pruned += prune_with_stats(&mut μ);
+        }
+
         let iter = state.iteration();
         stats.this_iters += 1;
 
         // Give statistics if needed
         state.if_verbose(|| {
-            plotter.plot_spikes(iter, maybe_d.as_ref(), Some(&τv), &μ);
+            plotter.plot_spikes(iter, None, Some(&τv), &μ);
             full_stats(&μ, ε, std::mem::replace(&mut stats, IterInfo::new()))
         });
 
@@ -292,14 +299,12 @@
 
     // Run the algorithm
     for state in iterator.iter_init(|| full_stats(&μ, ε, stats.clone())) {
+        let μ_base_len = μ.len();
         // Calculate smooth part of surrogate model.
         let mut τv = f.differential(&μ) * τ;
 
-        let μ_base_len = μ.len();
-
-        // Insert new spikes and reweigh
-        let (maybe_d, _within_tolerances) = prox_penalty
-            .insert_and_reweigh(&mut μ, &mut τv, τ, ε, config, &reg, &state, &mut stats)?;
+        // Do spike insertiona nd finite-dimensional weight optimisation
+        prox_penalty.insert_and_reweigh(&mut μ, &mut τv, τ, ε, config, &reg, &state, &mut stats)?;
 
         stats.inserted += μ.len() - μ_base_len;
 
@@ -310,6 +315,12 @@
             warned_merging = true;
         }
 
+        // Do extra weight optimisation step heuristic
+        for _ in 0..config.extra_weight_optimisation_steps {
+            τv = f.differential(&μ) * τ;
+            prox_penalty.reweigh(&mut μ, &mut τv, τ, ε, config, &reg, &state, &mut stats)?;
+        }
+
         // Update inertial prameters
         let λ_prev = λ;
         λ = 2.0 * λ_prev / (λ_prev + (4.0 + λ_prev * λ_prev).sqrt());
@@ -334,7 +345,7 @@
 
         // Give statistics if needed
         state.if_verbose(|| {
-            plotter.plot_spikes(iter, maybe_d.as_ref(), Some(&τv), &μ_prev);
+            plotter.plot_spikes(iter, None, Some(&τv), &μ_prev);
             full_stats(&μ_prev, ε, std::mem::replace(&mut stats, IterInfo::new()))
         });