Skip to content

Optimize MLX LRR training - #24

Merged
cboulay merged 1 commit into
devfrom
perf/mlx-lrr-training
Aug 24, 2026
Merged

Optimize MLX LRR training#24
cboulay merged 1 commit into
devfrom
perf/mlx-lrr-training

Conversation

@cboulay

@cboulay cboulay commented Aug 24, 2026

Copy link
Copy Markdown
Member

Summary

  • make zero-ridge MLX fitting robust to lazy inverse failures by selecting CPU-stream pinv up front, while retaining the faster inverse for positive ridge
  • asynchronously materialize the accumulated MLX covariance so incremental training does not retain every input chunk
  • assemble contiguous channel-group weights once instead of scattering each block through selection-matrix matmuls
  • add an opt-out for the synchronizing finite-value guard on trusted streams
  • repair the MLX training benchmark so it evaluates the actual covariance, weights, and effective transform

CCA and RefitKalman MLX compatibility are intentionally out of scope.

Profiling

M4 Pro, MLX 0.32.0, 512 channels in 8 groups of 64:

  • contiguous solve assembly: 1.520 ms to 1.396 ms on MLX; 2.329 ms to 0.629 ms on NumPy
  • 100 incremental updates: peak retained MLX memory 6.25 MiB to 0.17 MiB with async covariance evaluation
  • corrected end-to-end MLX training benchmark at 100 samples: 3.27 ms with zero-ridge pinv; 1.64 ms with ridge 1e-3 and inv

Validation

  • ruff check
  • ruff format --check
  • pytest: 360 passed, 24 skipped

@cboulay
cboulay merged commit 24534bd into dev Aug 24, 2026
9 checks passed
@cboulay
cboulay deleted the perf/mlx-lrr-training branch August 24, 2026 02:47
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.

1 participant