Oxidize ML Models

Machine learning models for Oxidize — a workforce access security platform with an agentic closed-loop relay runtime built in pure Rust.

Oxidize runs autonomous ML agents inline on the relay datapath. Models predict, decide, and act in real-time — from preemptive congestion control to per-hostname threat scoring to adaptive per-flow strategy selection.

Model Tiers

Tier 1 — Core Datapath Intelligence

Models that run on every packet/flow decision in the hot path.

Model	Architecture	Latency	File	Purpose
transformer_loss	MiniTransformer (d=64, 4 heads, INT8)	<10µs	`transformer_loss.safetensors`	Predict packet loss 50–100ms ahead for preemptive FEC injection
ppo_congestion	PPO Actor-Critic (continuous action, 128-unit hidden)	<1µs	`ppo_congestion.safetensors`	Smooth congestion window optimization — continuous multiplier, no CWND oscillation
security_threat_scorer	INT8 MLP (16 features → 64 hidden → 1)	<1µs	Inline (compiled into binary)	Per-hostname threat probability for browser isolation + SWG gating

Tier 2 — Per-Flow Optimization

Bandit and heuristic engines that adapt strategy per traffic class.

Model	Architecture	Latency	Purpose
flow_strategy	UCB1 bandit over 18 curated arms	<1µs	Select optimal (FEC level, compression, batching) per DPI-classified traffic class
compression_oracle	Entropy heuristic + magic-byte detection	<1µs	Decide Skip/Light/Normal/Aggressive compression per payload
path_selector	UCB1 contextual bandit (4 paths)	<1µs	Learn optimal path per traffic type across Primary/Secondary/Tertiary/Quaternary

Tier 3 — Prediction-Action Agent

Closed-loop autonomous agent that converts passive predictions into preemptive actions.

Component	What it does
PredictionActionAgent	State machine (Stable → PredictedDegradation → ActiveMitigation → Recovery) that emits `AgentAction`s when degradation probability exceeds threshold
AgentActions	PreemptiveFecBoost, CwndTaper, PathSwitch, HandoffPrep — each with outcome tracking that feeds reward back into PPO training
Autotune	Self-adjusts thresholds (FEC boost, CWND reduce, action cooldown) based on observed improvement rate

Tier 4 — Server-Side Agentic Engines

Relay-side engines that operate on connection metadata, not raw packets.

Engine	Purpose
AgentEgressEngine	Credential detection, OAuth/API-key/JWT/cloud-credential identification in egress traffic; verdict = Allow/Log/Alert/Block
AgentMemoryService	Per-entity risk memory with behavioral baselines, false-positive tracking, and time-decayed risk profiles
MobileThreatDefenseEngine	Mobile app risk scoring — jailbreak/root detection, app reputation, OS version risk, sideload detection
AgenticRecommendationsEngine	Generates prioritized security recommendations (policy tightening, access reviews, config hardening) from observed traffic patterns
HandoffPredictor	WiFi/LTE/5G signal tracking → predicts network handoff probability → triggers preemptive tunnel prep

Architecture Details

Transformer Loss Predictor (INT8 Quantized)

Input: [batch, 20, 8]  →  MultiHeadAttention(d=64, h=4)  →  FFN  →  Linear(1)  →  Sigmoid  →  Loss probability
                                    ↓
                        INT8 Quantization (10x speedup)
                                    ↓
                    Speculative Cache (next 100 decisions pre-computed)

Network Features (8):

Feature	Description	Normalization
`rtt_us`	Round-trip time	/ 1,000,000
`rtt_var_us`	RTT variance (jitter)	/ 500,000
`bandwidth_bps`	Estimated bandwidth	/ 10 Gbps
`loss_rate`	Recent packet loss	Raw (0–1)
`loss_trend`	Loss rate derivative	Raw
`inflight`	Packets in flight	/ 10,000
`cwnd`	Congestion window	/ 1,000,000
`buffer_occupancy`	Buffer fill level	Raw (0–1)

PPO Congestion Controller

Input: [batch, 8]  →  Actor(128)  →  ReLU  →  Actor(128)  →  ReLU  →  Gaussian(mean, log_std)  →  CWND multiplier
                      Critic(128)  →  ReLU  →  Critic(128)  →  ReLU  →  Linear(1)  →  State value

Continuous action space — outputs a CWND multiplier instead of discrete up/down, eliminating sawtooth oscillation.

Security Threat Scorer

INT8 MLP that scores hostnames for CompositeThreatIntel, which combines feed-based (URLhaus, OpenPhish, abuse.ch, PhishTank) and ML-based classification.

Threat Scoring Features (16 inputs):

Feature	Description	Normalization
`domain_entropy`	Shannon entropy of hostname characters	0–1
`subdomain_count_norm`	Domain length / 253	0–1
`suspicious_pattern`	Punycode, login-, secure-, excessive digits	0 or 1
`high_risk_tld`	.tk, .top, .xyz, etc.	0 or 1
`novelty_score`	First-seen recency	0–1
`distinct_actors_norm`	Distinct actors hitting host / 10	0–1
`ioc_match`	Request rate (requests/sec / 100)	0–1
`dlp_violation`	Payload entropy	0–1
`policy_violation`	SWG/CASB violation flag	0 or 1
`geo_risk`	Destination geo risk	0–1
`protocol_risk`	HTTPS=0, HTTP=0.5, raw TCP=1	0–1
`port_risk`	Well-known=0, registered=0.5, ephemeral=1	0–1

Threat classes: Malware, Phishing, C2, Botnet, Cryptojacking, Adware, DGA, Suspicious, Sensitive, Clean.

Provider architecture:

FeedProvider (O(1) hash lookups, 200+ built-in indicators)
  └── URLhaus, OpenPhish, abuse.ch, PhishTank
          │
          ▼
CompositeThreatIntel (Feed → ML → Cache)
  ├── Feed priority: confirmed IOCs checked first
  ├── ML fallback: INT8 MLP scores features per hostname
  └── DashMap TTL cache (32K cap)
          │
          ├──► BrowserIsolationEngine
          ├──► SwgEngine
          └──► Forwarder (ml_threat_score)

Flow Strategy Engine (DPI-Informed Bandit)

Per-traffic-class UCB1 bandit over 18 curated (FEC, Compression, Batching) arms:

Traffic Class	Reward Signal	Optimization Target
Gaming	1 / jitter_ms	Lowest latency, lowest jitter
VoIP	1 / latency_ms	Low RTT, stable connection
Streaming	1 / throughput_stddev	Highest bandwidth stability
Interactive	1 / latency_ms	Responsive, low-loss
Bulk	throughput_mbps (normalized)	Maximum throughput

Strategy dimensions:

FEC Level: None, Light, Medium, Heavy, Maximum
Compression: None, LZ4, Zstd, ROHC-only
Batching: Immediate, MicroBatch (250µs flush), StandardBatch (1ms flush)

Prediction-Action Agent State Machine

Stable ──(P(degradation) > threshold)──► PredictedDegradation
  ▲                                            │
  │                                    (emit AgentAction)
  │                                            ▼
Recovery ◄──(metrics improving)──── ActiveMitigation
  │                                            │
  └──(sustained good metrics)──────────────────┘

Actions and their triggers:

PreemptiveFecBoost — loss probability exceeds FEC threshold → boost redundancy before loss hits
CwndTaper — congestion predicted → gradual CWND reduction instead of cliff drop
PathSwitch — primary path degrading → switch to best alternative path
HandoffPrep — WiFi/LTE signal declining → pre-establish tunnel on new interface

Each action outcome (Improved/NoEffect/Degraded) feeds back as reward into PPO training via generate_training_feedback().

Advanced ML Features

Implemented in oxidize_common::advanced_ml:

Feature	What it does
FederatedCoordinator	Privacy-preserving distributed training across relay servers. Differential privacy (ε=1.0), secure aggregation, trust-weighted updates. Aggregates weight deltas every 60s when `min_clients` threshold met.
MultiAgentCoordinator	Per-flow cooperative congestion control via RL. Each connection = an agent. Inter-agent messages distributed every 10s for cooperative optimization.
ABTestingFramework	Welch's t-test based statistical significance for model variant deployment. Create experiments at runtime, stats logged in 30s summary tick.

Data Quality Guards

Built-in validation (ml_data_quality.rs) prevents training on garbage data:

Range validation — RTT [100µs, 10s], bandwidth [1Kbps, 100Gbps], loss rate [0, 1]
Consistency checks — RTT variance ≤ RTT, buffer occupancy ∈ [0, 1]
Timestamp validation — rejects future and stale samples
Anomaly detection — rejects duplicate/synthetic patterns (DDoS poisoning defense)
Threat sample validation — all 16 features ∈ [0, 1], no NaN/Inf, label ∈ [0, 1]

Self-healing: anomaly detection tracks validation failure rate; triggers automatic model retraining when >50% of samples are rejected.

Heuristic Fallback

If models aren't loaded, all components fall back to fast heuristics:

Component	Fallback
Loss Predictor	Exponential weighted moving average
Congestion Controller	Standard congestion avoidance
Compression Oracle	Entropy threshold + magic byte detection
Path Selector	Round-robin with availability check
Flow Strategy	Default strategy per traffic class
Threat Scorer	Feed-only lookup (no ML scoring)

Performance

All ML operations designed for minimal hot-path latency:

Operation	Typical Latency
Loss prediction (cached)	<1µs
Loss prediction (uncached)	<50µs
CWND optimization	<10µs
Path selection	<5µs
Flow strategy selection	<1µs
Threat scoring	<1µs
FEC decision	<100ns
Data quality validation	<1µs

Training runs asynchronously via BackgroundTrainer — never blocks the packet path.

Usage

With Oxidize (Rust)

use oxidize_common::ml_optimized::OptimizedMlEngine;

let engine = OptimizedMlEngine::new();

// Loss prediction (INT8 Transformer)
let loss_prob = engine.predict_loss(seq_num, &features);

// CWND optimization (PPO continuous)
let cwnd = engine.get_cwnd(rtt_us, &state);

// Threat scoring (INT8 MLP)
let threat_score = engine.predict_threat(&hostname_features);

// Compression decision (entropy heuristic)
let decision = engine.compression_decision(&payload);

// FEC decision with metrics recording
let fec = engine.fec_decision_with_metrics(&net_features, &metrics);

Flow Strategy Engine

use oxidize_common::flow_strategy::FlowStrategyEngine;

let mut engine = FlowStrategyEngine::new();

// Get optimal strategy for traffic class
let strategy = engine.get_strategy(TrafficClass::Gaming);
// strategy.fec_level, strategy.compression, strategy.batching

// Feed back observed quality
let reward = FlowStrategyEngine::gaming_reward(jitter_ms, loss_rate);
engine.record_reward(TrafficClass::Gaming, &strategy, reward);

Prediction-Action Agent

use oxidize_common::prediction_agent::PredictionActionAgent;

let mut agent = PredictionActionAgent::new();

// Evaluate current state → get preemptive actions
let actions = agent.evaluate(&ml_engine, &network_features);
for action in actions {
    match action {
        AgentAction::PreemptiveFecBoost { .. } => { /* boost FEC */ }
        AgentAction::CwndTaper { .. } => { /* reduce CWND */ }
        AgentAction::PathSwitch { .. } => { /* switch path */ }
        AgentAction::HandoffPrep { .. } => { /* prep handoff */ }
    }
}

// Record outcome for learning
agent.record_outcome(ActionOutcome::Improved);

Auto-Download Models

use oxidize_common::model_hub::ModelHub;

let hub = ModelHub::new(Default::default());
hub.download_models("/tmp/oxidize_models").await?;
engine.load_models("/tmp/oxidize_models")?;

Environment Variables

export HF_TOKEN=hf_xxxxxxxxxx          # For private repos or uploads
export OXIDIZE_MODEL_DIR=/tmp/oxidize_models  # Local cache

Training Pipeline

┌──────────────┐    ┌──────────────┐    ┌──────────────┐    ┌────────────┐
│   Servers    │───▶│   HF Hub     │───▶│   CI/CD      │───▶│  HF Hub    │
│  (collect)   │    │  (storage)   │    │  (train)     │    │  (models)  │
└──────────────┘    └──────────────┘    └──────────────┘    └────────────┘
       │                    ▲                   │                    │
       └────────────────────┴───────────────────┴────────────────────┘
                            Continuous Loop

Servers auto-collect — OptimizedMlEngine collects LossSample, DrlExperience, and ThreatSample training data
Servers auto-upload — Training data pushed to HF Hub hourly
CI trains daily — GitHub Actions aggregates and trains (3 AM UTC)
Servers auto-download — New models fetched from HF Hub on startup

Model Files

Model	File	Format
Loss Predictor	`transformer_loss.safetensors`	SafeTensors
Congestion Controller	`ppo_congestion.safetensors`	SafeTensors
Security Threat Scorer	Inline INT8 weights	Compiled into `oxidize-common` binary

Source Code Map

File	Contains
`common/src/ml_optimized.rs`	`OptimizedMlEngine`, `MiniTransformer`, `PPOController`, `SecurityScorer`, `SpeculativeCache`, `QuantizedLinear`
`common/src/prediction_agent.rs`	`PredictionActionAgent`, `AgentAction`, `AgentState`, autotune
`common/src/flow_strategy.rs`	`FlowStrategyEngine`, UCB1 bandit, 18 curated strategy arms
`common/src/threat_intel.rs`	`CompositeThreatIntel`, `FeedProvider`, `MlProvider`, threat class taxonomy
`common/src/advanced_ml.rs`	`FederatedCoordinator`, `MultiAgentCoordinator`, `ABTestingFramework`
`common/src/ml_data_quality.rs`	`DataQualityValidator`, validation thresholds, anomaly detection
`common/src/model_hub.rs`	`ModelHub`, HF Hub sync, `ModelConfig`
`common/src/handoff_prediction.rs`	`HandoffPredictor`, WiFi/LTE signal tracking
`server/src/agent_egress.rs`	`AgentEgressEngine`, credential detection, egress verdicts
`server/src/agent_memory.rs`	`AgentMemoryService`, per-entity risk memory, behavioral baselines
`server/src/mobile_threat_defense.rs`	`MobileThreatDefenseEngine`, mobile app risk scoring
`server/src/agentic_recommendations.rs`	`AgenticRecommendationsEngine`, policy/config recommendations

License

Citation

@software{oxidize2026,
  author = {gagansuie},
  title = {Oxidize ML Models},
  url = {https://github.com/gagansuie/oxidize},
  year = {2026}
}

gagansuie
/

oxidize-models