T-SKM-Net: A Trainable Neural Network Constraint Solving Framework Based on Sampled Kaczmarz-Motzkin Method

T-SKM-Net is a research result accepted by AAAI 2026, which transforms the classic Kaczmarz-Motzkin iterative method into a differentiable neural network layer, providing an end-to-end trainable solution for linear constraint satisfaction problems. This framework supports batch projection computation of mixed inequality and equality constraints, can be seamlessly integrated into deep learning frameworks like PyTorch, and is suitable for large-scale parallel computing and constraint-aware scenarios.

Linear constraint satisfaction problems are fundamental in the field of optimization, widely applied in machine learning, control engineering, and other scenarios. Traditional methods (such as interior-point methods and simplex methods) have efficiency bottlenecks in large-scale parallelism and end-to-end learning. The classic Kaczmarz method (for solving linear systems) and Motzkin method (for handling inequality constraints) are used independently, making it difficult to directly embed them into neural network training processes.

T-SKM-Net is a trainable sampled Kaczmarz-Motzkin network framework that supports batch computation of mixed linear constraint sets (A x ≤ b inequality constraints, C x = d equality constraints). The framework provides four algorithm variants:

1. Basic SKM: Randomly samples constraint subsets for projection, balancing convergence and computational complexity;
2. Heavy Ball SKM: Introduces a momentum term to accelerate convergence;
3. Nesterov Accelerated SKM: Uses a look-ahead step to improve convergence rate;
4. Generalized SKM: Supports custom constraint selection probability distributions.

The framework exposes two core hyperparameters:

• beta (β): Controls the number of sampled constraints, default is 6; a larger value increases coverage but also computational overhead;
• delta (δ): Step size scaling factor, default is 1.0, which needs to be adjusted based on the condition number of the problem. Both can be learned end-to-end via PyTorch optimizers to achieve adaptive configuration.

Application scenarios of T-SKM-Net include:

1. Differentiable constraint layer: Ensures network outputs meet physical/business rules, suitable for PINN and safety-critical applications;
2. Large-scale batch solving: GPU parallel computing, with significant speedup compared to traditional methods;
3. End-to-end learning: Embedded into neural network forward propagation, supports backpropagation training, providing tools for constraint-aware models.

The project uses a Python toolchain and can be quickly installed via uv (PyTorch needs to be installed first to adapt to the hardware). The code follows tensor shape conventions to ensure GPU efficiency: the inequality matrix A and equality matrix C are shared, and the initial point x0 and equality right-hand side d support batched inputs, balancing memory efficiency and flexibility.

This work was published in AAAI 2026 (Volume 40, Issue 17, DOI: 10.1609/aaai.v40i17.38459), with arXiv preprint and PDF versions available. For researchers in optimization algorithms and deep learning systems, T-SKM-Net is a tool worth exploring, and it is expected to promote the widespread application of differentiable optimization layers in industry.