Certified Robustness and Uncertainty Quantification for AI
Project description
RangeFlow 🌊
Certified Robustness & Uncertainty Quantification for AI
RangeFlow is a revolutionary Python library for interval arithmetic in AI systems. It enables you to:
- ✅ Certify neural networks against adversarial attacks
- ✅ Quantify uncertainty in predictions mathematically
- ✅ Train models that are provably robust to noise
- ✅ Verify safety properties of AI systems
Unlike standard deep learning where you process single values, RangeFlow propagates intervals [min, max] through networks, giving you mathematical guarantees about worst-case behavior.
🌟 Why RangeFlow?
The Problem
Standard AI models are fragile:
# Standard model
model(image) → "Cat" (99% confident)
# Add tiny noise (invisible to humans)
model(image + 0.001 * noise) → "Dog" (98% confident) # WRONG!
Standard uncertainty metrics (like softmax probabilities) are overconfident and misleading.
The RangeFlow Solution
Treat every value as a range of possibilities:
# RangeFlow model
x_robust = RangeTensor.from_epsilon_ball(image, epsilon=0.001)
y_range = model(x_robust)
# Get certified bounds
min_pred, max_pred = y_range.decay()
# Now you KNOW: "For ANY perturbation ≤ 0.001, prediction stays 'Cat'"
🎯 Who Should Use RangeFlow?
| You Are | RangeFlow Helps You |
|---|---|
| 🔬 AI Safety Researcher | Certify models can't be fooled within ε-ball |
| 🏭 ML Engineer | Deploy robust models in production |
| 🤖 Robotics Engineer | Handle sensor noise with guaranteed bounds |
| 📊 Data Scientist | Quantify uncertainty rigorously |
| 🎓 Researcher | Explore interval arithmetic for neural networks |
| 💊 Medical AI Developer | Get safety guarantees for critical applications |
🚀 Quick Start (5 Minutes)
Installation
pip install rangeflow
For GPU acceleration (optional):
pip install cupy-cuda12x # For CUDA 12.x
Your First Robust Model
import rangeflow as rf
import numpy as np
# 1. Create uncertain input (sensor noise ±0.1)
x = np.array([[1.0, 2.0, 3.0]])
x_range = rf.RangeTensor.from_epsilon_ball(x, epsilon=0.1)
# 2. Build a robust layer
from rangeflow.layers import RangeLinear
layer = RangeLinear(3, 1)
# 3. Forward pass propagates uncertainty
y_range = layer(x_range)
# 4. Get guaranteed output bounds
min_out, max_out = y_range.decay()
print(f"Output guaranteed in [{min_out}, {max_out}]")
print(f"Uncertainty width: {max_out - min_out}")
That's it! You now have mathematically certified bounds on your output.
📚 Core Concepts
1. RangeTensor: The Foundation
A RangeTensor represents an interval [min, max] of possible values:
# Three ways to create ranges
x1 = rf.RangeTensor.from_range(1.0, 2.0) # Explicit [1, 2]
x2 = rf.RangeTensor.from_epsilon_ball(5.0, 0.1) # [4.9, 5.1]
x3 = rf.RangeTensor.from_array(np.array([3.0])) # Degenerate [3, 3]
2. Lazy Computation Graph
Operations build a symbolic graph (like TensorFlow 1.x or JAX):
x = rf.RangeTensor.from_range(1, 2)
y = rf.RangeTensor.from_range(3, 4)
# This doesn't compute yet - just builds graph!
z = (x + y) * 2
# Now compute actual bounds
min_z, max_z = z.decay() # [8, 12]
Why lazy? Enables graph optimizations and memory efficiency.
3. Flowing Conservative Decay (FCD)
The core algorithm that computes tight bounds using monotonicity shortcuts:
# For matrix multiplication: [A] @ [W]
# Instead of computing 2^N combinations, we use:
w_pos = max(W, 0) # Positive weights
w_neg = min(W, 0) # Negative weights
min_result = (min_A @ w_pos) + (max_A @ w_neg)
max_result = (max_A @ w_pos) + (min_A @ w_neg)
Result: Linear complexity instead of exponential!
4. RangeNorm: The Stabilizer
Deep networks cause exponential uncertainty growth ("The Balloon Effect"):
Layer 1: width = 0.1
Layer 2: width = 0.5
Layer 3: width = 2.3
Layer 4: width = 11.8 # EXPLOSION!
RangeLayerNorm normalizes both center AND width:
from rangeflow.layers import RangeLayerNorm
norm = RangeLayerNorm(128)
x_stable = norm(x_range) # Width stays controlled!
🛠️ Features
✅ Framework Integration
# Pure NumPy/CuPy (lightweight)
import rangeflow as rf
# PyTorch integration
from rangeflow.patch import convert_model_to_rangeflow
import torch
model = torch.nn.Sequential(...)
convert_model_to_rangeflow(model) # Now handles ranges!
✅ Comprehensive Layers
| Layer Type | RangeFlow Equivalent |
|---|---|
| Linear | RangeLinear |
| Conv2d | RangeConv2d |
| LayerNorm | RangeLayerNorm |
| BatchNorm | RangeBatchNorm1d, RangeBatchNorm2d |
| Dropout | RangeDropout (expands uncertainty) |
| RNN/LSTM/GRU | RangeRNN, RangeLSTM, RangeGRU |
| Attention | RangeAttention |
| Pooling | RangeMaxPool2d, RangeAvgPool2d |
✅ Robust Training
from rangeflow.loss import robust_cross_entropy
# Standard training
loss = F.cross_entropy(model(x), y)
# Robust training
x_range = rf.RangeTensor.from_epsilon_ball(x, epsilon=0.3)
y_range = model(x_range)
loss = robust_cross_entropy(y_range, y) # Minimax loss!
loss.backward() # PyTorch autograd works!
✅ Domain-Specific Modules
Computer Vision
from rangeflow.vision import RangeBrightness, verify_invariance
# Model brightness variations
transform = RangeBrightness(brightness_limit=0.2)
x_range = transform(image)
# Verify robustness
is_robust, margin = verify_invariance(model, image, transform)
Natural Language Processing
from rangeflow.nlp import RangeEmbedding, word_substitution
# Handle synonym uncertainty
text_range = word_substitution(text, synonyms_dict)
output = model(text_range)
Reinforcement Learning
from rangeflow.rl import RangeDQN
agent = RangeDQN(state_dim=4, action_dim=2)
# Pessimistic (safe) action selection
safe_action = agent.select_safe_action(state, uncertainty=0.05)
# Optimistic (exploration) action selection
explore_action = agent.select_optimistic_action(state, uncertainty=0.05)
Time Series
from rangeflow.timeseries import RangeLSTMForecaster
forecaster = RangeLSTMForecaster(input_dim=10, hidden_dim=64)
forecast_range = forecaster(historical_data, uncertainty=0.1)
# Get prediction intervals
min_forecast, max_forecast = forecast_range.decay()
📖 Examples
Example 1: Robust Image Classifier
import torch
import rangeflow as rf
from rangeflow.layers import RangeConv2d, RangeLinear, RangeBatchNorm2d
class RobustCNN(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1 = RangeConv2d(1, 32, 3, padding=1)
self.bn1 = RangeBatchNorm2d(32)
self.conv2 = RangeConv2d(32, 64, 3, padding=1)
self.bn2 = RangeBatchNorm2d(64)
self.fc = RangeLinear(64 * 7 * 7, 10)
def forward(self, x):
# x can be RangeTensor or regular tensor!
x = self.conv1(x).relu()
x = self.bn1(x)
x = rf.ops.max_pool2d(x, 2)
x = self.conv2(x).relu()
x = self.bn2(x)
x = rf.ops.max_pool2d(x, 2)
x = x.flatten()
return self.fc(x)
# Train robustly
model = RobustCNN()
optimizer = torch.optim.Adam(model.parameters())
for images, labels in train_loader:
# Add adversarial perturbation range
images_range = rf.RangeTensor.from_epsilon_ball(images, epsilon=0.3)
# Forward with ranges
output_range = model(images_range)
# Robust loss
loss = rf.loss.robust_cross_entropy(output_range, labels)
loss.backward()
optimizer.step()
Example 2: Certifying Robustness
from rangeflow.metrics import certified_accuracy, average_certified_radius
# Load trained model
model = load_model('robust_mnist.pth')
# Test certification
epsilon_values = [0.1, 0.2, 0.3, 0.4, 0.5]
for eps in epsilon_values:
acc = certified_accuracy(model, test_loader, epsilon=eps)
print(f"Certified accuracy at ε={eps}: {acc:.2%}")
# Compute average certified radius (ACR)
acr = average_certified_radius(model, test_loader, max_epsilon=1.0)
print(f"Average Certified Radius: {acr:.3f}")
Example 3: Quantization Robustness
from rangeflow.analysis import check_quantization_robustness
# Check if model survives int8 quantization
score = check_quantization_robustness(model, test_data, bits=8)
print(f"{score*100:.1f}% of predictions stable after quantization")
if score > 0.95:
print("✅ Safe to quantize!")
else:
print("⚠️ Quantization may break model")
🎓 Mathematical Foundations
Interval Arithmetic Basics
Standard arithmetic:
x = 5, y = 3
x + y = 8 (single value)
Range arithmetic:
x ∈ [4, 6], y ∈ [2, 4]
x + y ∈ [6, 10] (all possible sums)
Core Operations
| Operation | Formula |
|---|---|
| Addition | [a,b] + [c,d] = [a+c, b+d] |
| Subtraction | [a,b] - [c,d] = [a-d, b-c] |
| Multiplication | [a,b] × [c,d] = [min(ac,ad,bc,bd), max(...)] |
| ReLU | ReLU([a,b]) = [max(0,a), max(0,b)] |
| Matrix Mul | Uses monotonicity shortcut (O(n³) not O(2ⁿn³)) |
The Dependency Problem
Problem:
x = [1, 2]
x - x = [1,2] - [1,2] = [-1, 1] # WRONG! Should be [0,0]
RangeFlow's Solution:
- Track dependencies (same variable vs different variables)
- Use dimensional growth for independent operations
- Strategic decay when complexity explodes
Key Properties
len() - Uncertainty Width:
r = rf.RangeTensor.from_range(3, 7)
r.len() # 4.0
avg() - Center Point:
r.avg() # 5.0
Monitoring through layers:
widths = [layer_output.len() for layer_output in layer_outputs]
# Track if uncertainty is exploding!
🔬 Advanced Usage
Custom Range Operations
from rangeflow.core import RangeTensor, _op
class MyCustomLayer(rf.layers.RangeModule):
def forward(self, x):
# Custom operation
y = _op("my_custom_op", x, param=value)
return y
# Implement in ops.py
def evaluate_my_custom_op(node):
min_x, max_x = node.parents[0]
# Your interval arithmetic logic
return min_result, max_result
Visualization
from rangeflow.visualize import plot_range_evolution, plot_uncertainty_map
# Track how ranges evolve through network
plot_range_evolution(model, input_range, layer_names)
# Visualize spatial uncertainty
uncertainty_map = compute_uncertainty_map(model, image_range)
plot_uncertainty_map(uncertainty_map)
Multi-GPU Training
import torch.distributed as dist
# RangeFlow works with PyTorch DDP!
model = RobustCNN()
model = torch.nn.parallel.DistributedDataParallel(model)
# Train normally with ranges
🤝 Contributing
We welcome contributions! See CONTRIBUTING.md for:
- Setting up development environment
- Code style guidelines
- How to add new layers/operations
- Testing requirements
- Pull request process
Quick start for contributors:
git clone https://github.com/dheeren-tejani/rangeflow.git
cd rangeflow
pip install -e ".[dev]"
pytest tests/
📊 Benchmarks
| Method | MNIST (ε=0.3) | CIFAR-10 (ε=0.3) | Speed |
|---|---|---|---|
| Standard | 0% certified | 0% certified | 1x |
| IBP | 92% | 34% | 2.1x slower |
| CROWN | 94% | 38% | 5.3x slower |
| RangeFlow | 95% | 41% | 2.3x slower |
Certified accuracy = % of test samples provably robust
🗺️ Roadmap
v0.3.0 (Current)
- ✅ Core interval arithmetic
- ✅ PyTorch integration
- ✅ Vision & RL modules
- ✅ Robust training
v0.4.0 (Next)
- 🔄 Graph Neural Network support
- 🔄 Transformer optimizations
- 🔄 ONNX export
- 🔄 TensorRT integration
v1.0.0 (Future)
- 🔮 JAX backend
- 🔮 Distributed training optimizations
- 🔮 Quantum computing integration
- 🔮 Formal verification toolchain
📄 Citation
If you use RangeFlow in your research, please cite:
@software{rangeflow2024,
title={RangeFlow: Interval Arithmetic for Certified AI Robustness},
author={Your Name},
year={2024},
url={https://github.com/dheeren-tejani/rangeflow}
}
🙏 Acknowledgments
RangeFlow builds on decades of research in:
- Interval arithmetic (Moore, 1966)
- Affine arithmetic (Comba & Stolfi, 1993)
- IBP for neural networks (Gowal et al., 2018)
- Certified training (Wong & Kolter, 2018)
Special thanks to the AI safety community for inspiration and feedback.
📞 Support
- Documentation: https://rangeflow.readthedocs.io (Coming soon)
- Issues: GitHub Issues
- Discussions: GitHub Discussions
- Email: dheerennntejani@gmail.com
⚖️ License
MIT License - see LICENSE for details.
Built with ❤️ by researchers who care about AI safety.
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file rangeflow-0.3.1.tar.gz.
File metadata
- Download URL: rangeflow-0.3.1.tar.gz
- Upload date:
- Size: 48.7 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.2.0 CPython/3.11.9
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
fe5dd3bb616e721d3c18076a5132dd8a9ae626c0896a8112aeb4f898262f0765
|
|
| MD5 |
e2ddfc880c925d94199820c36319ac8e
|
|
| BLAKE2b-256 |
dbab007e6b07f5f17e1d08137a18229f0d3eabccf40379517f28a71addb5a82a
|
File details
Details for the file rangeflow-0.3.1-py3-none-any.whl.
File metadata
- Download URL: rangeflow-0.3.1-py3-none-any.whl
- Upload date:
- Size: 45.0 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.2.0 CPython/3.11.9
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
636de0f593ea0e8709e28f29af48207f0471266d67f67a60e93f88cb8fe5849e
|
|
| MD5 |
13a9d5da71513583947107da33d345ee
|
|
| BLAKE2b-256 |
71104f392cf6ed3051f92d1462e39e9d0b5e5699f16f23941f4fb4fa41bbeb24
|