Skip to content

li-group/DisjunctiveNet.jl

BranchesTags

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

20 Commits
examples
examples
 
 
src
src
 
 
test
test
 
 
.gitignore
.gitignore
 
 
LICENSE
LICENSE
 
 
Project.toml
Project.toml
 
 
README.md
README.md
 
 

Repository files navigation

DisjunctiveNet.jl

DisjunctiveNet.jl provides differentiable projection layers for neural networks whose outputs must satisfy mixed logical-linear constraints.

The package lets users define constraints of the form

global linear constraints
+
(disjunct 1 OR disjunct 2 OR ...)
+
(disjunct 1 OR disjunct 2 OR ...)
+ ...

and automatically builds a differentiable projection layer that can be placed after a neural network, especially a Flux.jl model.

The intended workflow is:

raw neural prediction
        ↓
differentiable disjunctive projection layer
        ↓
constraint-satisfying prediction

The package is designed for research on differentiable constrained learning, neural-symbolic learning, and safe neural network prediction.

Formulations

The layer supports three lifted convexified formulations.

Formulation Description
:dnf Full DNF convex-hull formulation. Stronger but can grow combinatorially.
:cnf One convex-hull block per disjunction. Scalable but generally weaker.
:partial_dnf Expands a chosen subset of rules into DNF and keeps the rest as CNF.

All three formulations are convexified projection layers. The projected output is guaranteed to satisfy the selected convexified formulation. Depending on the formulation, especially :cnf and :partial_dnf, the projected output may lie in a relaxation of the original discrete logical set.

Installation

This package currently uses a modified version of DiffOpt.jl with a regularized KKT solve for more stable reverse differentiation.

Until the DiffOpt patch is upstreamed, install the modified DiffOpt first:

using Pkg

Pkg.add(url = "https://github.com/ShramanPal/DiffOpt.jl")
Pkg.add(url = "https://github.com/ShramanPal/DisjunctiveNet.jl")

For local development:

git clone https://github.com/ShramanPal/DisjunctiveNet.jl
cd DisjunctiveNet.jl
julia --project=.

Then:

using Pkg
Pkg.instantiate()
Pkg.test()

Basic idea

A neural network produces an unconstrained prediction:

yhat = backbone(x)

The projection layer maps it to a point satisfying the selected convexified rule set:

y = projection_layer(yhat)

The full constrained model is:

x -> backbone(x) -> disjunctive projection layer -> projected y

The projection is differentiable, so training can be performed using ordinary Flux/Zygote code.

A first disjunctive model

Suppose the neural network predicts two outputs:

y[1], y[2]

We want:

0 <= y[1] <= 1
0 <= y[2] <= 1

global constraint:
y[1] + y[2] >= 0.8

rule 1:
y[1] <= 0.25 OR y[1] >= 0.75

rule 2:
y[2] <= 0.25 OR y[2] >= 0.75

This can be modeled as follows.

using DisjunctiveDifferentiableLayers

# Create a disjunctive model with two output variables:
# y[1] and y[2].
dm = DisjunctiveModel(2)

# Set box bounds:
# 0 <= y[1] <= 1
# 0 <= y[2] <= 1
set_bounds!(
    dm,
    lower = [0.0, 0.0],
    upper = [1.0, 1.0],
)

# Add a global linear constraint.
# This constraint must hold in the projected output:
# y[1] + y[2] >= 0.8
add_linear_constraint!(dm, [1.0, 1.0], :>=, 0.8)

# Add rule 1:
# y[1] <= 0.25 OR y[1] >= 0.75
#
# The name :x_split is optional, but useful for partial-DNF ordering.
add_disjunction!(
    dm,
    [LinearConstraint([1.0, 0.0], :<=, 0.25)],
    [LinearConstraint([1.0, 0.0], :>=, 0.75)];
    name = :x_split,
)

# Add rule 2:
# y[2] <= 0.25 OR y[2] >= 0.75
#
# The name :y_split lets us refer to this rule later, for example in
# rule_ordering = [:x_split, :y_split].
add_disjunction!(
    dm,
    [LinearConstraint([0.0, 1.0], :<=, 0.25)],
    [LinearConstraint([0.0, 1.0], :>=, 0.75)];
    name = :y_split,
)

Inspect the user-facing model:

print_model(dm)

Inspect the lifted projection formulations:

print_projection_model(dm; formulation = :dnf)

print_projection_model(dm; formulation = :cnf)

print_projection_model(
    dm;
    formulation = :partial_dnf,
    num_dnf_rules = 1,
    rule_ordering = [:x_split, :y_split],
)

Optional: JuMP-style linear constraints

For convenience, users can also write constraints using JuMP’s @build_constraint syntax. This avoids manually constructing LinearConstraint objects.

using JuMP
using DisjunctiveDifferentiableLayers

# Create a fresh model.
dm2 = DisjunctiveModel(2)

set_bounds!(
    dm2,
    lower = [0.0, 0.0],
    upper = [1.0, 1.0],
)

# Get the JuMP variables associated with the model outputs.
y = output_variables(dm2)

# Add the global constraint:
# y[1] + y[2] >= 0.8
add_linear_constraint!(
    dm2,
    @build_constraint(y[1] + y[2] >= 0.8),
)

# Add rule 1:
# y[1] <= 0.25 OR y[1] >= 0.75
add_disjunction!(
    dm2,
    [@build_constraint(y[1] <= 0.25)],
    [@build_constraint(y[1] >= 0.75)];
    name = :x_split,
)

# Add rule 2:
# y[2] <= 0.25 OR y[2] >= 0.75
add_disjunction!(
    dm2,
    [@build_constraint(y[2] <= 0.25)],
    [@build_constraint(y[2] >= 0.75)];
    name = :y_split,
)

The explicit LinearConstraint(...) interface and the JuMP-style @build_constraint(...) interface create the same internal representation.

Core API

DisjunctiveModel(n_outputs)

Creates a model for a neural network output vector of length n_outputs.

dm = DisjunctiveModel(3)

This means the projection layer expects raw neural network predictions of length 3:

yhat = [yhat1, yhat2, yhat3]

and returns a projected vector:

y = [y1, y2, y3]

set_bounds!(model; lower, upper)

Sets box bounds on the projected output variables.

set_bounds!(
    dm,
    lower = [0.0, 0.0],
    upper = [1.0, 1.0],
)

This imposes:

0 <= y[1] <= 1
0 <= y[2] <= 1

Finite bounds are strongly recommended because they are used by the convex-hull formulations.

LinearConstraint(a, sense, b)

Represents a linear constraint of the form:

a' y <= b
a' y >= b
a' y == b

depending on sense.

Examples:

LinearConstraint([1.0, 1.0], :>=, 0.8)

means:

y[1] + y[2] >= 0.8

LinearConstraint([1.0, 0.0], :<=, 0.25)

means:

y[1] <= 0.25

Supported senses are:

:<=
:>=
:(==)

add_linear_constraint!(model, ...)

Adds a global linear constraint.

Global constraints are enforced across the full projection model. They are useful for constraints that should always hold, regardless of which disjunctive branch is active.

Coefficient-vector form:

add_linear_constraint!(
    dm,
    [1.0, 1.0],
    :>=,
    0.8,
)

Equivalent explicit object form:

add_linear_constraint!(
    dm,
    LinearConstraint([1.0, 1.0], :>=, 0.8),
)

Equivalent JuMP-style form:

using JuMP

y = output_variables(dm)

add_linear_constraint!(
    dm,
    @build_constraint(y[1] + y[2] >= 0.8),
)

add_disjunction!(model, disjuncts...; name = nothing)

Adds a disjunctive rule.

Each argument is one disjunct, and each disjunct is a vector of linear constraints.

For example:

add_disjunction!(
    dm,
    [LinearConstraint([1.0, 0.0], :<=, 0.25)],
    [LinearConstraint([1.0, 0.0], :>=, 0.75)];
    name = :x_split,
)

means:

y[1] <= 0.25 OR y[1] >= 0.75

A disjunct can contain multiple constraints:

add_disjunction!(
    dm,
    [
        LinearConstraint([1.0, 0.0], :>=, 0.4),
        LinearConstraint([1.0, 0.0], :<=, 0.6),
    ],
    [LinearConstraint([1.0, 0.0], :>=, 0.8)];
    name = :x_mid_or_high,
)

means:

0.4 <= y[1] <= 0.6 OR y[1] >= 0.8

The same rule can be written with JuMP-style constraints:

using JuMP

y = output_variables(dm)

add_disjunction!(
    dm,
    [
        @build_constraint(y[1] >= 0.4),
        @build_constraint(y[1] <= 0.6),
    ],
    [@build_constraint(y[1] >= 0.8)];
    name = :x_mid_or_high,
)

The name keyword is optional, but recommended. Named rules can be used to control partial-DNF ordering:

rule_ordering = [:x_split, :y_split]

output_variables(model)

Returns JuMP variables associated with the output coordinates. This is used with JuMP-style constraint construction.

using JuMP

y = output_variables(dm)

add_linear_constraint!(
    dm,
    @build_constraint(y[1] + y[2] >= 0.8),
)

This is often easier to read than manually writing coefficient vectors.

project(model, yhat; formulation = :dnf)

Projects a raw prediction onto the selected convexified rule set.

yhat = [0.5, 0.1]

result = project(dm, yhat; formulation = :dnf)

println(result.status)
println(result.y)

The result contains:

result.y       # projected prediction
result.status  # solver termination status
result.model   # internal JuMP/DiffOpt model

Common choices:

project(dm, yhat; formulation = :dnf)

project(dm, yhat; formulation = :cnf)

project(
    dm,
    yhat;
    formulation = :partial_dnf,
    num_dnf_rules = 1,
    rule_ordering = [:x_split, :y_split],
)

DisjunctiveProjectionLayer(model; formulation = :dnf)

Creates a differentiable projection layer.

layer = DisjunctiveProjectionLayer(dm; formulation = :partial_dnf)

y = layer(yhat)

The layer can be placed after a Flux neural network.

model = Chain(
    Dense(4 => 16, relu),
    Dense(16 => 2),
    layer,
)

constrained_model(backbone, model; kwargs...)

Wraps a Flux neural network with a disjunctive projection layer.

model = constrained_model(
    backbone,
    dm;
    formulation = :partial_dnf,
    num_dnf_rules = 1,
    rule_ordering = [:x_rule, :y_rule],
)

Calling the model evaluates:

yhat = backbone(x)
y = projection_layer(yhat)

Only the neural network backbone has trainable parameters. The projection layer is differentiable but has no trainable weights.

constrained_model(backbone, example_input) do dm ... end

Creates a DisjunctiveModel automatically by inferring the output dimension from the neural network.

model = constrained_model(backbone, x0; formulation = :cnf) do dm
    set_bounds!(dm, lower = zeros(2), upper = ones(2))
    add_linear_constraint!(dm, [1.0, 1.0], :>=, 0.8)
end

This is the most convenient API for Flux users.

Projection API

Project a raw prediction:

yhat = [0.5, 0.1]

result = project(dm, yhat; formulation = :dnf)

println(result.status)
println(result.y)

Use a projection layer directly:

layer = DisjunctiveProjectionLayer(dm; formulation = :dnf)

y = layer(yhat)

Available formulations:

layer_dnf = DisjunctiveProjectionLayer(dm; formulation = :dnf)

layer_cnf = DisjunctiveProjectionLayer(dm; formulation = :cnf)

layer_partial = DisjunctiveProjectionLayer(
    dm;
    formulation = :partial_dnf,
    num_dnf_rules = 1,
    rule_ordering = [:x_split, :y_split],
)

Full Flux example

The following example builds a small neural network, defines a disjunctive rule set, compares :dnf, :cnf, and :partial_dnf projections on the same raw neural network output, runs one training step, and performs inference on a single sample.

The same script is available under examples/flux_end2end.jl.

using Flux
using Zygote
using DisjunctiveDifferentiableLayers

# -----------------------------
# 1. Build a neural network
# -----------------------------

backbone = Chain(
    Dense(3 => 8, relu),
    Dense(8 => 2),
)

x = Float32[0.2, 0.7, 0.4]
target = Float32[0.8, 0.2]

# -----------------------------
# 2. Build the disjunctive rule model
# -----------------------------

dm = DisjunctiveModel(2)

set_bounds!(
    dm,
    lower = [0.0, 0.0],
    upper = [1.0, 1.0],
)

# Global constraint:
# y1 + y2 >= 0.8
add_linear_constraint!(dm, [1.0, 1.0], :>=, 0.8)

# Rule 1 has three disjuncts:
# y1 <= 0.2 OR 0.4 <= y1 <= 0.6 OR y1 >= 0.8
add_disjunction!(
    dm,
    [LinearConstraint([1.0, 0.0], :<=, 0.2)],
    [
        LinearConstraint([1.0, 0.0], :>=, 0.4),
        LinearConstraint([1.0, 0.0], :<=, 0.6),
    ],
    [LinearConstraint([1.0, 0.0], :>=, 0.8)];
    name = :x_rule,
)

# Rule 2 has three disjuncts:
# y2 <= 0.2 OR 0.35 <= y2 <= 0.55 OR y2 >= 0.7
add_disjunction!(
    dm,
    [LinearConstraint([0.0, 1.0], :<=, 0.2)],
    [
        LinearConstraint([0.0, 1.0], :>=, 0.35),
        LinearConstraint([0.0, 1.0], :<=, 0.55),
    ],
    [LinearConstraint([0.0, 1.0], :>=, 0.7)];
    name = :y_rule,
)

# -----------------------------
# 3. Inspect the model and formulations
# -----------------------------

println()
println("=== User-facing disjunctive model ===")
print_model(dm)

println()
println("=== DNF lifted formulation ===")
print_projection_model(dm; formulation = :dnf)

println()
println("=== CNF lifted formulation ===")
print_projection_model(dm; formulation = :cnf)

println()
println("=== Partial-DNF lifted formulation ===")
print_projection_model(
    dm;
    formulation = :partial_dnf,
    num_dnf_rules = 1,
    rule_ordering = [:x_rule, :y_rule],
)

# -----------------------------
# 4. Compare DNF, CNF, and partial-DNF on the same NN output
# -----------------------------

yhat = backbone(x)

println()
println("Raw NN output yhat = ", yhat)

dnf_layer = DisjunctiveProjectionLayer(
    dm;
    formulation = :dnf,
    y_regularization = 1e-4,
    ycopy_regularization = 1e-4,
    gamma_regularization = 1e-4,
    anchor_regularization = 1e-4,
)

cnf_layer = DisjunctiveProjectionLayer(
    dm;
    formulation = :cnf,
    y_regularization = 1e-4,
    ycopy_regularization = 1e-4,
    gamma_regularization = 1e-4,
    anchor_regularization = 1e-4,
)

partial_layer = DisjunctiveProjectionLayer(
    dm;
    formulation = :partial_dnf,
    num_dnf_rules = 1,
    rule_ordering = [:x_rule, :y_rule],
    y_regularization = 1e-4,
    ycopy_regularization = 1e-4,
    gamma_regularization = 1e-4,
    anchor_regularization = 1e-4,
)

y_dnf = dnf_layer(yhat)
y_cnf = cnf_layer(yhat)
y_partial = partial_layer(yhat)

println()
println("=== Projected predictions for the same sample ===")
println("DNF projection         = ", y_dnf, "  sum = ", sum(y_dnf))
println("CNF projection         = ", y_cnf, "  sum = ", sum(y_cnf))
println("Partial-DNF projection = ", y_partial, "  sum = ", sum(y_partial))

println()
println("The three projections can differ because :dnf, :cnf, and :partial_dnf")
println("construct different convexified relaxations of the same logical rule set.")

# -----------------------------
# 5. Build a trainable constrained Flux model
# -----------------------------

# Here we train the partial-DNF constrained model.
# The projection layer is differentiable, but it has no trainable parameters.
# Flux only trains the neural network backbone.
model = constrained_model(
    backbone,
    dm;
    formulation = :partial_dnf,
    num_dnf_rules = 1,
    rule_ordering = [:x_rule, :y_rule],
    y_regularization = 1e-4,
    ycopy_regularization = 1e-4,
    gamma_regularization = 1e-4,
    anchor_regularization = 1e-4,
)

# -----------------------------
# 6. Forward pass
# -----------------------------

y = model(x)

println()
println("=== Forward pass through constrained model ===")
println("Projected prediction = ", y)
println("Feasibility check: y1 + y2 = ", sum(y))

# -----------------------------
# 7. One training step
# -----------------------------

loss(m, x, target) = sum(abs2, m(x) .- target)

opt = Flux.setup(Adam(1e-3), model)

l, grads = Flux.withgradient(model) do m
    loss(m, x, target)
end

Flux.update!(opt, model, grads[1])

println()
println("Training loss before update = ", l)

# -----------------------------
# 8. Inference after one update
# -----------------------------

y_after = model(x)

println()
println("=== Inference after one update ===")
println("Projected prediction after one update = ", y_after)
println("Feasibility check after update: y1 + y2 = ", sum(y_after))

The constrained model is a Flux-compatible model.

Formulation summaries and benchmarking

The package includes utilities for inspecting lifted formulations:

formulation_summary(dm; formulation = :dnf)

formulation_summary(dm; formulation = :cnf)

formulation_summary(
    dm;
    formulation = :partial_dnf,
    num_dnf_rules = 1,
    rule_ordering = [:x_split, :y_split],
)

Benchmark projection time and model size:

benchmark_projection(dm, yhat; formulation = :cnf, label = "CNF")

benchmark_projection(
    dm,
    yhat;
    formulation = :partial_dnf,
    num_dnf_rules = 1,
    rule_ordering = [:x_split, :y_split],
    label = "partial-DNF k=1",
)

Example output:

CNF                  status=OPTIMAL build=0.002s solve=0.003s vars=54 cons=152 y=[0.5, 0.2]
partial-DNF k=1      status=OPTIMAL build=0.003s solve=0.004s vars=62 cons=180 y=[0.5, 0.2]

Examples

Run examples from the package root:

julia --project=. examples/basic_example.jl
julia --project=. examples/midsize_example.jl
julia --project=. examples/largesize_example.jl
julia --project=. examples/flux_end2end.jl

Current limitations

  • Constraints are currently linear.
  • The differentiable backend relies on DiffOpt.
  • A modified DiffOpt with regularized KKT solves is recommended for robust reverse differentiation.
  • Full DNF can grow exponentially in the number of disjunctions.
  • Partial-DNF is intended to trade off formulation strength and computational size.
  • The package currently targets Flux.jl as the main neural network interface.

Citation

If you use this package in your research, please cite:

@inproceedings{
anonymous2026disjunctivenet,
title={DisjunctiveNet: Neural Symbolic Learning via Differentiable Convexified Optimization Layers},
author={Anonymous},
booktitle={Forty-third International Conference on Machine Learning},
year={2026},
url={https://openreview.net/forum?id=c88GPpURN8}
}

Once the paper is de-anonymized, the author field will be updated.

License

This package is released under the MIT License. See LICENSE for details.

About

Disjunctive differentiable layers julia package

Resources

Readme

License

MIT license
Activity
Custom properties

Stars

1 star

Watchers

0 watching

Forks

0 forks
Report repository

Releases

No releases published

Packages

 
 
 

Contributors

Languages