Skip to content

Latest commit

 

History

History
300 lines (213 loc) · 6.93 KB

File metadata and controls

300 lines (213 loc) · 6.93 KB

The Axiom.jl Vision

"What if your ML model couldn't have bugs?"


The Crisis in Machine Learning

We have a problem. A big one.

Machine learning models are being deployed in:

  • Medical diagnosis - Wrong predictions can kill
  • Autonomous vehicles - Wrong predictions can kill
  • Financial systems - Wrong predictions can destroy lives
  • Criminal justice - Wrong predictions can imprison innocents

And yet, our tools for building these systems are... inadequate.

The State of ML Engineering

# This is how we build AI systems that make life-or-death decisions
class MedicalAI(nn.Module):
    def __init__(self):
        self.conv1 = nn.Conv2d(3, 64, 3)
        self.fc1 = nn.Linear(64 * 30 * 30, 10)  # Probably right?

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = x.view(x.size(0), -1)  # Reshape... hope this is correct
        x = self.fc1(x)  # RuntimeError after 3 hours of training
        return F.softmax(x, dim=1)

# Did you catch the bug? There are actually 3.

This is insane.

We would never accept this in aerospace. In nuclear power. In bridge construction.

Why do we accept it in AI?


The Axiom Thesis

What if ML frameworks worked like compilers?

Compilers catch bugs before your code runs. They verify types, check syntax, ensure consistency.

Axiom.jl brings this to machine learning:

@axiom MedicalAI begin
    input :: Tensor{Float32, (224, 224, 3)}
    output :: Probabilities(10)

    features = input |> Conv(64, (3,3))
    output = features |> Dense(10)  # COMPILE ERROR!

    # "Shape mismatch at Dense layer
    #  Expected input: Vector
    #  Got: Tensor{Float32, (222, 222, 64)}
    #
    #  Solution: Add Flatten layer between Conv and Dense
    #
    #  Would you like me to fix this? [y/N]"
end

The bug is caught immediately. Not after training. Not in production. Now.


Three Pillars of Axiom.jl

1. Compile-Time Verification

# In Axiom.jl, the type system encodes tensor shapes
input :: Tensor{Float32, (batch, 28, 28, 1)}

# Shape errors are COMPILE errors
Dense(10)(input)  # Error: Dense expects 2D, got 4D

# Correct code:
Flatten()(input) |> Dense(784, 10)  # ✓ Compiles

What this means: You can't run code with shape errors. Period.

2. Formal Guarantees

@axiom Classifier begin
    input :: Image
    output :: Probabilities(10)

    # ... layers ...

    # @ensure adds runtime contracts; @prove attempts to discharge them statically
    @ensure sum(output) ≈ 1.0
    @ensure all(output .>= 0)
    @prove ∀x. no_nan(output(x))
end

What this means: @ensure attaches runtime contracts that are checked on every forward pass; @prove attempts to discharge a property statically — via known-pattern heuristics, or an SMT solver when SMTLib.jl is loaded — and honestly returns :unknown when it cannot. It is not a blanket claim that every property is formally proved.

3. Production Performance

# Development: Julia backend (fast iteration)
model = compile(MyModel, backend=JuliaBackend())

# Production: Zig backend (native SIMD kernels)
model = compile(MyModel, backend=ZigBackend("/path/to/libaxiom_zig.so"), optimize=:aggressive)
# Competitive with PyTorch on small/medium workloads — see benchmark/ for measured medians

What this means: verification guarantees without giving up competitive performance.


Why Julia + Zig?

Why Not Pure Python?

Python is dynamically typed. You can't encode shape constraints in the type system.

# This is valid Python - no way to prevent it
def broken(x):
    return torch.matmul(x, torch.randn(999, 999))  # Probably wrong

# Only crashes at runtime, maybe

Why Not a Pure Systems Language?

Systems languages like Zig are great for native kernels. But ML research needs:

  • REPL exploration - Try ideas instantly
  • Interactive visualization - Plot results immediately
  • Rapid iteration - Change code, see results
// Edit code
// Wait 30s-2min for compilation
// Run
// Find bug
// Repeat

// This is a flow killer for research

Why Julia + Zig?

Julia for research:

julia> model = Sequential(Dense(10, 5), ReLU())
julia> model(randn(4, 10))  # Instant feedback
4×5 Matrix{Float32}:
 0.0  0.123  0.0  0.456  0.789
 ...

Zig for production:

# When you're done experimenting
production_model = compile(model, backend=ZigBackend("/path/to/libaxiom_zig.so"))
# Native SIMD kernels; competitive on small/medium ops (see benchmark/)

Best of both worlds.


The Verification Ladder

Axiom.jl provides multiple levels of verification:

Level 1: Shape Checking (Automatic)

# The compiler catches this
Conv(64, (3,3))(input_2d)  # Error: Expected 4D input

# No code to write - it just works

Level 2: Runtime Assertions (@ensure)

@axiom Model begin
    # ...
    @ensure sum(output) ≈ 1.0  # Checked at runtime
    @ensure all(output .>= 0)
end

Level 3: Formal Proofs (@prove)

@axiom Model begin
    # ...
    @prove ∀x. sum(softmax(x)) == 1.0  # discharged via @prove (known pattern; SMT when SMTLib.jl loaded)
    @prove ∀x ε. (ε < δ) ⟹ stable(f(x), f(x+ε))  # Robustness
end

Level 4: Certification

# Generate proof certificate for regulatory approval
cert = generate_certificate(model, properties)
save_certificate(cert, "fda_submission.cert")

Who Is This For?

ML Researchers

  • Catch bugs faster
  • Iterate more quickly
  • Publish reproducible results

ML Engineers

  • Deploy with confidence
  • Debug production issues
  • Meet performance requirements

Safety-Critical Applications

  • Medical AI (FDA approval)
  • Autonomous vehicles (safety certification)
  • Financial systems (regulatory compliance)
  • Aerospace (DO-178C compliance)

The Road Ahead

Phase 1: Foundation (Current)

  • Core DSL and type system
  • Basic verification (@ensure)
  • Julia backend
  • PyTorch import

Phase 2: Performance

  • Full Rust backend
  • GPU acceleration (CUDA, Metal)
  • Distributed training
  • ONNX export

Phase 3: Verification

  • SMT solver integration
  • Automated proof generation
  • Robustness certification
  • Fairness verification

Phase 4: Ecosystem

  • Model zoo (verified models)
  • Hugging Face integration
  • Cloud deployment
  • Industry certifications

Join the Revolution

We're building the future of machine learning. A future where:

  • Bugs are caught before they cause harm
  • Models come with mathematical guarantees
  • Safety and performance aren't trade-offs
  • AI systems can be trusted

Want to help?


"The best way to predict the future is to invent it." — Alan Kay

Let's invent verified ML together.