- A high-performance Consensus ADMM optimization toolkit designed for
- solving distributed convex problems in large-scale machine learning and numerical optimization tasks.
- GitHub Repository URL:
- https://github.com/Allenwang2004/consensus-admm
- The library provides:
- Low-level C++ core with fast linear algebra routines, parallel updates,
- and memory-efficient structures.
Because of the increasing size of datasets and models, many optimization problems in machine learning and statistics are naturally distributed across multiple nodes. Consensus ADMM is a powerful algorithm to solve such problems efficiently by decomposing them into smaller subproblems that can be solved in parallel while enforcing consensus among the variables.
- These problems typically appear in:
- Distributed Lasso / ElasticNet regression: split data across machines,
- reach shared model.
- Federated Learning: multiple clients optimize local models while
- maintaining global consistency.
- PCA / Matrix Completion: decompose large matrices with consensus
- constraints.
- Input
- Objective functions f_i(x_i) encoded as callable Python functions.
- Optional proximal operators for non-smooth terms (e.g., L1, indicator).
- Consensus constraints (e.g., x_i = z) implicitly handled.
- Output
- Converged primal variables x_i, consensus estimate z, dual variables u_i,
- and convergence history.
- Logs / profiling / residuals for debugging and diagnostics.
- Constraints
- Problem must be convex and separable (i.e., sum f_i(x_i) + consensus
- constraints).
- Supports smooth + non-smooth objectives.
- Current version supports synchronous updates (async under consideration).
Python Interface
from consensus_admm import ConsensusADMM
def local_loss_i(x): ...
def proximal_g_i(x, rho): ...
solver = ConsensusADMM(
num_agents=10,
local_losses=[local_loss_1, ..., local_loss_10],
local_prox_ops=[prox_g1, ..., prox_g10],
rho=1.0,
max_iters=1000,
tol=1e-4
)
results = solver.solve()
z = results['consensus']
x_list = results['locals']
history = results['residuals']
Main Class
class ConsensusADMM:
def __init__(
self,
num_agents: int,
local_losses: list[Callable[[np.ndarray], float]],
local_prox_ops: list[Callable[[np.ndarray, float], np.ndarray]],
rho: float = 1.0,
max_iters: int = 1000,
tol: float = 1e-4
):
"""Initialize ADMM solver with agent-wise loss and proximal
updates."""
def solve(self) -> dict:
"""Run ADMM iterations and return consensus + local solutions."""import numpy as np
from consensus_admm import ConsensusADMM, soft_threshold
# Example: Distributed Lasso regression
num_agents = 3
variable_dim = 10
# Create loss functions for each agent
def make_quadratic_loss(A, b):
return lambda x: 0.5 * np.linalg.norm(A @ x - b) ** 2
def make_quadratic_grad(A, b):
return lambda x: A.T @ (A @ x - b)
# L1 proximal operator for sparsity
lambda_reg = 0.1
def l1_prox(x, rho):
return soft_threshold(x, lambda_reg / rho)
# Generate synthetic data
np.random.seed(42)
data = [(np.random.randn(5, variable_dim), np.random.randn(5))
for _ in range(num_agents)]
# Set up solver
solver = ConsensusADMM(
num_agents=num_agents,
local_losses=[make_quadratic_loss(A, b) for A, b in data],
local_gradients=[make_quadratic_grad(A, b) for A, b in data],
local_prox_ops=[l1_prox] * num_agents,
rho=1.0,
max_iters=1000,
tol=1e-4,
verbose=True
)
# Solve
results = solver.solve(
initial_x=[np.zeros(variable_dim)] * num_agents,
initial_z=np.zeros(variable_dim)
)
print(f"Converged: {results['converged']}")
print(f"Consensus solution: {results['consensus']}")- Build System
- CMake for compiling the C++ backend
- pybind11 for C++ ↔ Python bindings
- setup.py for Python package installation
- Optional: Support for CUDA backend in future versions
- Licensing
- Apache 2.0 License
- Testing Framework
- C++ Unit Tests with GoogleTest:
- Numerical correctness of updates
- Memory management and edge cases
- Python Tests using pytest:
- API compliance and regression tests
- Distributed optimization examples (Lasso, Ridge)
- Documentation
- README.rst: install, usage, examples, FAQ
- Python: rich docstrings + type hints
- C++: Doxygen-style comments
Week Milestone
- 09/27 Repository setup, literature review, project skeleton(C++ core + Python binding)
- 10/04 Implement primal/dual updates and support for L1, L2 proximal operators
- 10/11 Benchmark on synthetic distributed Lasso && Ridge regression and python binding with Pybind11
- 10/18 Add plotting / convergence diagnostics and prepare for PCA / Matrix Completion knowledge
- 10/25 Extend to PCA / Matrix Completion via ADMM formulation
- 11/01 Optimize performance, memory usage; profile bottlenecks
- 11/08 Finalize README, write docs
- 11/15 Stretch goal: multi-threading
- Boyd, S., Parikh, N., Chu, E., Peleato, B., & Eckstein, J. (2010).
- Distributed Optimization and Statistical Learning via the Alternating Direction Method of Multipliers
- pybind11 Documentation: https://pybind11.readthedocs.io/
- CMake Documentation: https://cmake.org/documentation/
- Eigen C++ Linear Algebra: https://eigen.tuxfamily.org
This project is licensed under the MIT License - see the LICENSE file for details.
Copyright (c) 2025 Allenwang2004
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.