---
language:
  - en
license: mit
tags:
  - secret-detection
  - mixture-of-experts
  - gating-network
  - security
  - nlp
  - token-classification
pipeline_tag: token-classification
library_name: pytorch
datasets:
  - custom
metrics:
  - accuracy
  - f1
model-index:
  - name: secretmask-gate
    results:
      - task:
          type: routing
          name: Expert Routing
        dataset:
          name: SecretMask v2
          type: custom
        metrics:
          - type: accuracy
            value: 0.927
            name: Test Accuracy
          - type: accuracy
            value: 1.0
            name: Validation Accuracy
base_model:
  - andrewandrewsen/distilbert-secret-masker
  - andrewandrewsen/longformer-secret-masker
---

# SecretMask MoE Gating Network

[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
[![Model Size](https://img.shields.io/badge/Size-12KB-green.svg)](https://huggingface.co/andrewandrewsen/secretmask-gate)

**Lightweight learned gating network for SecretMask Mixture-of-Experts routing.**

This repository contains a trained 12KB neural network that learns optimal routing between two secret detection expert models. Use this for true MoE inference with weighted expert combination.

---

## 📋 Overview

The gating network is a tiny 3-layer MLP (3,042 parameters) that:

1. Takes 10 features extracted from text
2. Outputs routing weights `[w_fast, w_long]` (sum to 1.0)
3. Enables weighted combination of expert model outputs

**Training Results:**

- ✅ 100% validation accuracy (200 examples)
- ✅ 92.7% test accuracy (600 examples)
- ✅ Only 0.19ms inference overhead
- ✅ Matches heuristic routing performance

> **Note**: This gating network is **optional and experimental**. Heuristic (rule-based) routing achieves identical results (92.7% accuracy) without requiring this model. The recommended production configuration uses **Fast Expert + Filters** without learned routing or the Long Expert. This gate is primarily for learning/experimentation with MoE architectures. See [Configuration Guide](https://github.com/AndrewAndrewsen/secmask/blob/main/CONFIGURATION_GUIDE.md) for details.

---

## 🚀 Quick Start

### Installation

```bash
pip install torch transformers huggingface-hub
```

### Download and Use

```python
from huggingface_hub import hf_hub_download
from moe_gate import GatingNetwork, extract_features_tensor

# Download gating network
gate_path = hf_hub_download("andrewandrewsen/secretmask-gate", "best_gate.pt")

# Load model
gate = GatingNetwork.load(gate_path)
gate.eval()

# Extract features from text
text = "AWS key: AKIAIOSFODNN7EXAMPLE"
features = extract_features_tensor(text)

# Get routing weights
import torch
with torch.no_grad():
    weights = gate(features.unsqueeze(0))

print(f"Fast expert weight: {weights[0][0]:.3f}")
print(f"Long expert weight: {weights[0][1]:.3f}")
# Output: Fast expert weight: 0.950, Long expert weight: 0.050
```

### Integration with SecretMask

```bash
# Clone SecretMask repository
git clone https://github.com/andrewandrewsen/secmask.git
cd secmask

# Run inference with learned MoE routing
python infer_moe.py \
    --text "My AWS key is AKIAIOSFODNN7EXAMPLE" \
    --routing-mode learned \
    --fast-model andrewandrewsen/distilbert-secret-masker \
    --long-model andrewandrewsen/longformer-secret-masker \
    --gate-model andrewandrewsen/secretmask-gate \
    --tau 0.80
```

---

## 🏗️ Model Architecture

```
Input: [10 features]
    ↓
Linear(10 → 64) + LayerNorm + ReLU + Dropout(0.1)
    ↓
Linear(64 → 32) + LayerNorm + ReLU + Dropout(0.1)
    ↓
Linear(32 → 2) + Softmax
    ↓
Output: [w_fast, w_long]  (sum = 1.0)
```

**Total Parameters:** 3,042  
**Model Size:** 12KB (float32)  
**Inference Time:** ~0.19ms on CPU

---

## 📊 Input Features (10D)

The gating network takes a normalized 10-dimensional feature vector:

| Index | Feature           | Description             | Normalization |
| ----- | ----------------- | ----------------------- | ------------- |
| 0     | `token_count`     | Number of tokens        | / 1000        |
| 1     | `entropy`         | Shannon entropy         | / 6           |
| 2     | `has_pem`         | Has PEM block (binary)  | 0 or 1        |
| 3     | `has_k8s`         | Has K8s secret (binary) | 0 or 1        |
| 4     | `akia_count`      | AWS pattern count       | / 5           |
| 5     | `github_count`    | GitHub token count      | / 5           |
| 6     | `jwt_count`       | JWT token count         | / 5           |
| 7     | `base64_count`    | Base64 pattern count    | / 50          |
| 8     | `line_count`      | Number of lines         | / 100         |
| 9     | `avg_line_length` | Avg chars per line      | / 100         |

---

## 📈 Training Details

**Dataset:**

- Training: 6,000 examples
- Validation: 200 examples
- Test: 600 examples

**Configuration:**

- Optimizer: AdamW (lr=0.001, weight_decay=0.01)
- Scheduler: Cosine annealing
- Batch size: 32
- Epochs: 10
- Device: Apple M-series (MPS)

**Training Results:**

| Epoch | Train Loss | Train Acc | Val Loss | Val Acc  |
| ----- | ---------- | --------- | -------- | -------- |
| 1     | 0.0808     | 97.6%     | 0.0051   | **100%** |
| 2     | 0.0036     | 100%      | 0.0010   | **100%** |
| 10    | 0.0005     | 100%      | 0.0001   | **100%** |

**Test Performance:**

- Routing accuracy: 92.7%
- Fast expert: 92.7% of examples
- Long expert: 7.3% of examples
- Matches heuristic routing exactly

---

## 🔧 Usage with Expert Models

This gating network coordinates two expert models:

| Expert   | Model                                                                                                       | Size  | Max Tokens | Use Case                     |
| -------- | ----------------------------------------------------------------------------------------------------------- | ----- | ---------- | ---------------------------- |
| **Fast** | [andrewandrewsen/distilbert-secret-masker](https://huggingface.co/andrewandrewsen/distilbert-secret-masker) | 265MB | 512        | Short texts, code snippets   |
| **Long** | [andrewandrewsen/longformer-secret-masker](https://huggingface.co/andrewandrewsen/longformer-secret-masker) | 592MB | 2048       | Long documents, config files |

### How It Works

```python
# 1. Extract features
features = extract_features_tensor(text)

# 2. Get routing weights from gating network
weights = gate(features)  # [w_fast, w_long]

# 3. Run both expert models
fast_output = fast_expert(text)
long_output = long_expert(text)

# 4. Combine outputs using learned weights
final_output = weights[0] * fast_output + weights[1] * long_output
```

---

## 📦 Files in This Repository

- **`best_gate.pt`** - Trained gating network (12KB)
- **`final_gate.pt`** - Final checkpoint (12KB)
- **`history.json`** - Training history (3.2KB)
- **`README.md`** - This file

---

## 🔬 Technical Details

### Load Balancing

The model was trained with a load balancing loss to encourage uniform expert usage:

```python
target_distribution = [0.5, 0.5]  # 50% fast, 50% long
actual_distribution = weights.mean(dim=0)
load_balance_loss = 0.01 * MSE(actual_distribution, target_distribution)
```

Despite this, the model learned to route 90.5% to fast expert and 9.5% to long expert, matching the natural data distribution.

### Routing Metrics

```python
from moe_gate import compute_routing_metrics

weights = gate(features)
metrics = compute_routing_metrics(weights)

# Returns:
# {
#   'fast_expert_pct': 92.7,
#   'long_expert_pct': 7.3,
#   'avg_fast_weight': 0.924,
#   'avg_long_weight': 0.076,
#   'entropy': 0.031
# }
```

Low entropy (0.031) indicates confident routing decisions.

---

## 🆚 Heuristic vs Learned Routing

| Metric                | Heuristic            | Learned MoE              |
| --------------------- | -------------------- | ------------------------ |
| **Routing Accuracy**  | 92.7%                | 92.7%                    |
| **Model Size**        | 0KB (rules only)     | 12KB                     |
| **Latency**           | 0.065ms              | 0.256ms                  |
| **Training Required** | No                   | Yes (10 epochs)          |
| **Explainability**    | High (if-else rules) | Medium (learned weights) |
| **Adaptability**      | Manual updates       | Data-driven              |

**Recommendation:** Use heuristic routing for simplicity and explainability. Use learned routing when you want to fine-tune on your specific data distribution.

---

## 📚 Citation

If you use this model, please cite:

```bibtex
@model{secretmask-gate,
  author = {Anders Andersson},
  title = {SecretMask MoE Gating Network},
  year = {2025},
  publisher = {HuggingFace},
  url = {https://huggingface.co/andrewandrewsen/secretmask-gate}
}
```

---

## 📄 License

MIT License - see [LICENSE](LICENSE) file.

**Note:** This model is trained to work with the SecretMask expert models, which use Apache 2.0 licensed base models (DistilBERT, Longformer). See the expert model repositories for full licensing details.

---

## 🔗 Related Resources

- **SecretMask MoE Repository:** [GitHub](https://github.com/andrewandrewsen/secmask)
- **Fast Expert Model:** [andrewandrewsen/distilbert-secret-masker](https://huggingface.co/andrewandrewsen/distilbert-secret-masker)
- **Long Expert Model:** [andrewandrewsen/longformer-secret-masker](https://huggingface.co/andrewandrewsen/longformer-secret-masker)
- **Documentation:** See repository for BENCHMARKS.md, USE_CASES.md, etc.

---

## 🤝 Contributing

Issues and pull requests welcome at [GitHub](https://github.com/andrewandrewsen/secmask).

---

**Built with ❤️ for the open source community**