CVXPYlayers

repository·master·Indexed 24 days ago

https://github.com/cvxpy/cvxpylayers

A Python library for constructing differentiable convex optimization layers in PyTorch, JAX, and MLX using CVXPY. It solves parametrized convex optimization problems in the forward pass and computes derivatives in the backward pass. The library supports GPU acceleration via the Moreau solver and provides specialized examples for control systems, finance, machine learning, resource allocation, and physics.

Tokens
32.6K
Snippets
88
Records
126
Agent score
83%

What's inside cvxpylayers

  1. Key Features of CVXPYlayers

    master

    CVXPYlayers provides several advantages for deep learning research and production:

    • Encode Domain Knowledge: Inject physical constraints, safety bounds, or fairness requirements directly into your models by formulating them as convex programs.
    • GPU Acceleration: Solvers like Moreau and CuClarabel keep computations on the GPU, avoiding expensive CPU-GPU transfers.
    • Batched Solving: The first dimension of input tensors is treated as the batch dimension, allowing you to solve thousands of optimization problems in parallel.
    • Multiple Solvers: Choose from various solvers including Moreau, Clarabel, SCS, and CuClarabel depending on your problem structure.
  2. Optimization examples by domain

    master

    CVXPYlayers provides several specialized examples across different fields:

    Control Systems

    • Linear Quadratic Regulator (LQR): Learning optimal value function parameters.
    • Constrained LQR: LQR with control input bounds and state constraints.
    • Vehicle Path Planning: Autonomous vehicle trajectory optimization.
    • Model Predictive Control (MPC): MPC with learned cost-to-go functions.
    • Approximate Dynamic Programming: Convex approximations for dynamic programming.

    Finance & Portfolio Optimization

    • Markowitz Portfolio: Mean-variance optimization with dynamic rebalancing.
    • Portfolio with VIX: Volatility-aware optimization using the VIX index.

    Machine Learning

    • Monotonic Regression: Learning monotonic input-output relationships.
    • Signal Denoising: Signal/image denoising with learned parameters.
    • ReLU Layers: Optimization layers with ReLU activations.
    • Data Poisoning Attack: Adversarial attacks on machine learning models.

    Resource Allocation

    • Resource Allocation: Water and resource distribution optimization.
    • Supply Chain: Supply chain network flow optimization.

    Engineering

    • Stiffness Constants: Optimizing mechanical stiffness parameters.
    • GP Circuit Sizing: Geometric programming for transistor sizing (using gp=True).

    Optimization Techniques

    • Batched Portfolio with Duals: Batched solving, parameter broadcasting, and dual variable extraction.
    • Robust Optimization: Learning robustness from data for worst-case design.
    • Predict-then-Optimize: End-to-end pipelines combining neural networks and CvxpyLayer (Decision-focused learning).
    • Convex Regression: Fitting convex piecewise-linear functions via SGD.
  3. Overview of CVXPYlayers core concepts

    master

    CVXPYlayers converts CVXPY optimization problems into differentiable layers that can be integrated into machine learning models. To use it successfully, you must understand these four core concepts:

    1. DPP Compliance: Problems must follow Disciplined Parametrized Programming (DPP) rules to ensure they are valid for differentiation.
    2. Parameters vs Variables: In the context of a layer, Parameters act as the inputs to the optimization problem, while Variables act as the outputs.
    3. Implicit Differentiation: Gradients of the solution with respect to the parameters are computed using the KKT (Karush-Kuhn-Tucker) conditions.
    4. Cone Program Representation: CVXPY problems are canonicalized into a standard cone program form before being solved.
  4. Batching dual variables

    master

    Dual variables support batched parameters. When using batched tensors for parameters, the returned dual variables will also be batched. For a scalar constraint, the dual variable output will have a shape corresponding to the batch size.

    # Batch of 32 problems
    c_batch = torch.randn(32, n, requires_grad=True)
    b_batch = torch.randn(32, requires_grad=True)
    
    x_opt, eq_dual = layer(c_batch, b_batch)
    # x_opt: (32, n)
    # eq_dual: (32,) for scalar constraint
  5. Understand the CVXPYlayers data flow

    master

    The data flow within a CVXPYlayer follows a specific pipeline to transform input tensors into differentiable outputs:

    1. Parameters: Receives input tensors (supported frameworks include torch, jax, and mlx).
    2. Validate & Batch: Checks tensor shapes and handles broadcasting.
    3. Canonicalize: Converts the CVXPY problem into a standard cone program.
    4. Solve: Executes the optimization using solvers like diffcp, CuClarabel, etc.
    5. Extract Variables: Maps the numerical solution back to the original CVXPY variables.
    6. Variables: Returns the variables with gradients attached for backpropagation.
    Parameters (torch/jax/mlx tensors)
        |
        v
    +-------------------+
    | Validate & Batch  |
    +-------------------+
        |
        v
    +-------------------+
    | Canonicalize      |
    +-------------------+
        |
        v
    +-------------------+
    | Solve             |
    +-------------------+
        |
        v
    +-------------------+
    | Extract Variables |
    +-------------------+
        |
        v
    Variables (with gradients attached)
  6. Ensure your problem is DPP-compliant

    master

    For CvxpyLayer to work, the CVXPY problem must satisfy Disciplined Parametrized Programming (DPP) requirements:

    1. Parameters appear affinely in the objective and constraints.
    2. Problem structure is fixed (the set of constraints does not change based on parameter values).
    3. No parameter-dependent domains (the domain of variables cannot depend on parameters).

    Example of a valid (DPP-compliant) problem:

    import cvxpy as cp
    
    x = cp.Variable(2)
    A = cp.Parameter((3, 2))
    b = cp.Parameter(3)
    
    # Good: Parameter appears affinely
    problem = cp.Problem(cp.Minimize(cp.sum_squares(A @ x - b)))
    assert problem.is_dpp()  # True
    import cvxpy as cp
    
    x = cp.Variable(2)
    A = cp.Parameter((3, 2))
    b = cp.Parameter(3)
    
    # Good: Parameter appears affinely
    problem = cp.Problem(cp.Minimize(cp.sum_squares(A @ x - b)))
    assert problem.is_dpp()  # True
  7. Compare CvxpyLayer implementations across frameworks

    master

    While the constructor is consistent, the base class and the method used to execute the layer (the forward pass) depend on the framework you are using:

    • PyTorch: Inherits from torch.nn.Module and uses the forward(*params) method.
    • JAX: Implements a callable class and uses the __call__(*params) method.
    • MLX: Implements a callable class and uses the __call__(*params) method.
  8. Ensure problem compliance with DPP

    master

    Your CVXPY problem must follow Disciplined Parametrized Programming (DPP) rules. DPP ensures that the problem structure remains fixed while only the parameter values change, which is a requirement for implicit differentiation.

    To verify compliance, use problem.is_dpp().

    Examples:

    • Valid: A @ x - b (Parameters appear linearly).
    • Invalid: cp.quad_form(x, P) where P is a parameter (changes the problem structure).
    assert problem.is_dpp(), "Problem must be DPP-compliant"
  9. Use broadcasting with mixed batched and unbatched parameters

    master

    You can mix batched and unbatched parameters. Unbatched parameters are automatically broadcast across the batch dimension. This is useful when some parameters vary per instance while others are shared across the entire batch.

    # A is batched, b is shared across the batch
    A_batch = torch.randn(batch_size, 3, 2)  # Different A for each instance
    b_shared = torch.randn(3)                 # Same b for all instances
    
    (x_batch,) = layer(A_batch, b_shared)    # x_batch shape: (10, 2)
  10. How batching works in CVXPYlayers

    master

    CVXPYlayers supports solving multiple problem instances in parallel by adding a batch dimension as the first dimension of your parameter tensors.

    For a single instance with parameters of shape (D1, D2, ...), a batch of size N will have parameter shapes of (N, D1, D2, ...). The resulting variable tensors will also include this batch dimension at the front.

    import cvxpy as cp
    import torch
    from cvxpylayers.torch import CvxpyLayer
    
    # Problem with parameters of shape (3, 2) and (3,)
    x = cp.Variable(2)
    A = cp.Parameter((3, 2))
    b = cp.Parameter(3)
    problem = cp.Problem(cp.Minimize(cp.sum_squares(A @ x - b)), [x >= 0])
    
    layer = CvxpyLayer(problem, parameters=[A, b], variables=[x])
    
    # Single instance: shapes (3, 2) and (3,)
    A_single = torch.randn(3, 2)
    b_single = torch.randn(3)
    (x_single,) = layer(A_single, b_single)  # x_single shape: (2,)
    
    # Batched: shapes (batch_size, 3, 2) and (batch_size, 3)
    batch_size = 10
    A_batch = torch.randn(batch_size, 3, 2)
    b_batch = torch.randn(batch_size, 3)
    (x_batch,) = layer(A_batch, b_batch)  # x_batch shape: (10, 2)
  11. Getting Started with CVXPYlayers

    master

    To begin using cvxpylayers, select a quick start script or tutorial based on your preferred deep learning framework:

    PyTorch

    • Quick Start: torch/torch_example.py
    • Tutorial: torch/tutorial.ipynb

    JAX

    • Quick Start: jax/jax_example.py
    • Tutorial: jax/tutorial.ipynb

    MLX

    • Quick Start: mlx/mlx_example.py
  12. Install CVXPYlayers

    master

    You can install CVXPYlayers using pip. Choose the installation method that matches your needs:

    • Most Users (PyTorch): Includes PyTorch and all dependencies.
    • Minimal Install: Installs only the core package; you must add framework extras manually.
    • All Frameworks: Installs everything, useful for development or testing.