feat(aggregation): Add custom batched QP solver#596
Draft
PierreQuinton wants to merge 4 commits intomainfrom
Draft
feat(aggregation): Add custom batched QP solver#596PierreQuinton wants to merge 4 commits intomainfrom
PierreQuinton wants to merge 4 commits intomainfrom
Conversation
Replaces the previous QP solver with a pure-PyTorch ADMM solver for the projection-onto-dual-cone subproblem. Three techniques are combined: - **Ruiz equilibration** (10 iterations): symmetrically scales G so that every row/column has infinity-norm ≈ 1, reducing the effective condition number before factorization. - **ADMM**: splits the constrained QP into a cheap V-update (Cholesky solve of G_s + ρI) and a trivial Z-update (componentwise clamp onto the feasible set), following the OSQP formulation. - **Adaptive ρ**: every √m iterations, scales ρ up or down by 10× when primal and dual residuals are severely imbalanced, triggering a cheap re-factorization to keep convergence well-behaved. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
The ADMM solver in dual_cone.py always factorizes G_s + rho*I where rho > 0, so it handles positive semi-definite Gramians without needing an external regularization term. The 1e-4*I padding previously added in the forward methods is now redundant. Side-effect: the 1e-4*I pad was also acting as a strong preconditioner that tightened ADMM convergence across row permutations. Without it, permutation-invariance errors reflect the solver's actual accuracy on ill-conditioned inputs (~2e-5 for DualProj, ~1e-5 for UPGrad). Tolerances in the corresponding tests are updated accordingly. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
PierreQuinton
added
cc: feat
PierreQuinton
changed the title
github-actions
bot
changed the title
Contributor
|
Could we support both solvers so that we have one for small m and one for large m? |
Contributor
Author
|
Of course we can, ideally we want to explore the optimal choices and provide an automatic choice, but also provide a customizable one if needed. |
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
2 participants
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.
There is a compromise here, and it seems to be necessary. The algorithm used by
qpsolverhas great precision, but it cannot be run in parallel and does not compromise precision with time.This algorithm, has less precision, it is slower for some range of
m(medium, on CPU). But it is (most probably) much faster for large m on GPU. We could probably improve a bit more by implementing it as a cuda kernel.If we want uncompromising improvement here, we will never be able to change, we are probably at the Pareto front, maximizing precision.