README.md

---
license: mit
language:
- en
tags:
- sentence-transformers
- sentence-embeddings
- multi-task-learning
- reinforcement-learning
- semantic-similarity
- nli
- paraphrase-detection
datasets:
- sentence-transformers/stsb
- nyu-mll/multi_nli
- quora
- google-research-datasets/paws
- nyu-mll/glue
pipeline_tag: sentence-similarity
base_model: sentence-transformers/all-MiniLM-L6-v2
model-index:
- name: FireDevourerEmbedder-RL-v3.6
  results:
  - task:
      type: semantic-similarity
      name: Semantic Textual Similarity
    dataset:
      type: sentence-transformers/stsb
      name: STS-B
    metrics:
    - type: pearson_spearman_avg
      value: 0.3366
  - task:
      type: natural-language-inference
      name: Natural Language Inference
    dataset:
      type: nyu-mll/multi_nli
      name: MultiNLI
    metrics:
    - type: accuracy
      value: 0.7465
  - task:
      type: text-classification
      name: Question Duplicate Detection
    dataset:
      type: quora
      name: QQP
    metrics:
    - type: accuracy
      value: 0.8636
  - task:
      type: text-classification
      name: Paraphrase Detection
    dataset:
      type: google-research-datasets/paws
      name: PAWS
    metrics:
    - type: accuracy
      value: 0.8459
  - task:
      type: text-classification
      name: Paraphrase Detection
    dataset:
      type: nyu-mll/glue
      name: MRPC
    metrics:
    - type: accuracy
      value: 0.7744
---

# FireDevourerEmbedder-RL-v3.6

A multi-task sentence embedding model that uses **Reinforcement Learning** to dynamically optimize task weights during training. The model learns to balance multiple NLU tasks simultaneously, producing robust sentence embeddings suitable for semantic similarity, natural language inference, and paraphrase detection.

## Key Innovation

FireDevourerEmbedder introduces an **RL-based adaptive task weighting system** that automatically adjusts the importance of each training task based on validation performance. Instead of using fixed task weights, a policy network learns optimal weight distributions during training, leading to better overall performance across diverse NLU benchmarks.

## Why Multi-Task? Information-Dense Embeddings

The core philosophy behind FireDevourerEmbedder is that **multi-task learning creates richer, more information-dense embeddings** than single-task approaches.

By training with multiple task heads simultaneously, the shared encoder is forced to learn representations that capture:

| Dimension | Learned From | What It Captures |
|-----------|--------------|------------------|
| **Semantic Similarity** | STS-B | Fine-grained meaning overlap |
| **Logical Relationships** | MultiNLI | Entailment, contradiction, neutrality |
| **Question Semantics** | QQP | Intent and duplicate detection |
| **Adversarial Patterns** | PAWS | Word-order sensitivity, paraphrase robustness |
| **Domain Awareness** | All datasets | Context-appropriate representations |

This results in embeddings that are:
- **More robust** - trained to handle diverse linguistic phenomena
- **More transferable** - generalize better to unseen tasks
- **More informative** - each dimension of the embedding vector carries meaningful semantic signal

Unlike single-task embedders that optimize for one objective, FireDevourerEmbedder's embeddings simultaneously encode multiple facets of meaning, making them suitable for a wide range of downstream applications without fine-tuning.

## Model Details

| Property | Value |
|----------|-------|
| **Base Model** | [sentence-transformers/all-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2) |
| **Hidden Size** | 384 |
| **Version** | v3.6 |
| **Training Steps** | 80,000 |
| **Total Parameters** | ~22M |

## Architecture

The model consists of a shared BERT encoder with task-specific output heads:

```
Input Sentence(s)
       │
       ▼
┌─────────────────────────┐
│   MiniLM-L6-v2 Encoder  │
│     (384-dim output)    │
└─────────────────────────┘
       │
       ▼
   Mean Pooling
       │
       ├──► STS Head (384→1) ──► Similarity Score [0,1]
       ├──► NLI Head (384→3) ──► [Contradiction, Neutral, Entailment]
       ├──► QQP Head (384→2) ──► [Not Duplicate, Duplicate]
       ├──► PAWS Head (384→2) ──► [Not Paraphrase, Paraphrase]
       └──► Domain Head (384→5) ──► [General, Entailment, Questions, Adversarial, News]
```

## Performance

| Task | Dataset | Metric | Score |
|------|---------|--------|-------|
| Question Duplicate Detection | QQP | Accuracy + F1 | **0.8636** |
| Paraphrase Detection | PAWS | Accuracy + F1 | **0.8459** |
| Paraphrase Detection | MRPC | Accuracy + F1 | **0.7744** |
| Natural Language Inference | MultiNLI | Accuracy + F1 | **0.7465** |
| Semantic Textual Similarity | STS-B | Pearson/Spearman | **0.3366** |
| | | **Average** | **0.7134** |

## Training Details

### Datasets

The model was trained on 5 balanced datasets with 100,000 samples each (500,000 total):

| Dataset | Task Type | Domain | Samples |
|---------|-----------|--------|---------|
| [STS-B](https://huggingface.co/datasets/sentence-transformers/stsb) | Semantic Similarity | General | 100,000 |
| [MultiNLI](https://huggingface.co/datasets/nyu-mll/multi_nli) | Natural Language Inference | Entailment | 100,000 |
| [QQP](https://huggingface.co/datasets/quora) | Duplicate Question Detection | Questions | 100,000 |
| [PAWS](https://huggingface.co/datasets/google-research-datasets/paws) | Paraphrase Detection | Adversarial | 100,000 |
| [MRPC](https://huggingface.co/datasets/nyu-mll/glue) | Paraphrase Detection | News | 100,000 |

### Data Augmentation Strategy

To prevent training bias, all datasets were balanced to exactly **100,000 samples** each:

| Dataset | Original Size | Augmentation Method |
|---------|---------------|---------------------|
| STS-B | ~8,600 | Repetition (~12x) + pair swapping |
| MultiNLI | ~433,000 | Subsampling |
| QQP | ~400,000 | Subsampling |
| PAWS | ~49,000 | Repetition (~2x) + pair swapping |
| MRPC | ~3,600 | Repetition (~10x, capped) + pair swapping |

**Why this matters:**
- Without balancing, larger datasets (QQP, MultiNLI) would dominate training
- Smaller but valuable datasets (MRPC, STS-B) would be underrepresented
- Equal representation ensures the model learns equally from all task types

**Augmentation techniques:**
- **Repetition**: Smaller datasets repeated up to 10x maximum to prevent memorization
- **Sentence pair swapping**: For symmetric tasks, (A, B) pairs also trained as (B, A)

### Training Configuration

| Parameter | Value |
|-----------|-------|
| Epochs | 3 |
| Batch Size | 16 |
| Learning Rate | 2e-5 |
| Total Steps | 93,750 |
| Warmup Steps | 9,375 (10%) |
| Evaluation Frequency | Every 10,000 steps |
| Early Stopping | 3 consecutive decreases |
| Training Time | 3.29 hours |

### RL Weight Adaptation System

The model uses a policy network to dynamically adjust task weights during training:

| Parameter | Value |
|-----------|-------|
| RL Learning Rate | 0.001 |
| State Dimension | 6 (5 task scores + average) |
| Action Dimension | 5 (weight deltas) |
| Hidden Dimension | 32 |
| Delta Scale | ±5% per update |
| Update Frequency | Every 10,000 steps |

**Weight Evolution During Training:**

| Task | Initial Weight | Final Weight | Change |
|------|---------------|--------------|--------|
| STS | 0.250 | 0.282 | +0.032 |
| NLI | 0.300 | 0.337 | +0.037 |
| QQP | 0.200 | 0.063 | -0.137 |
| PAWS | 0.150 | 0.173 | +0.023 |
| MRPC | 0.100 | 0.145 | +0.045 |

The RL system learned to reduce QQP weight (already high-performing) while increasing weights for harder tasks.

## Training Progress

| Version | Step | Average Score | Best Task | Improvement |
|---------|------|---------------|-----------|-------------|
| v3.1 | 10,000 | 0.6133 | QQP (0.8093) | +0.6133 |
| v3.2 | 20,000 | 0.6430 | QQP (0.8351) | +0.0297 |
| v3.3 | 30,000 | 0.6813 | QQP (0.8391) | +0.0383 |
| v3.4 | 40,000 | 0.6925 | QQP (0.8527) | +0.0111 |
| v3.5 | 50,000 | 0.7099 | QQP (0.8579) | +0.0175 |
| **v3.6** | **80,000** | **0.7134** | **QQP (0.8636)** | **+0.0035** |

## Usage

### Installation

```bash
pip install torch transformers
```

### Loading the Model

```python
import torch
from transformers import AutoTokenizer, AutoModel

# Load tokenizer and base model
tokenizer = AutoTokenizer.from_pretrained("path/to/FireDevourerEmbedder-RL-v3.6")
base_model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")

# Load checkpoint
checkpoint = torch.load("path/to/FireDevourerEmbedder-RL-v3.6/full_checkpoint.pt")

# Load model weights (you'll need to reconstruct the full model class)
# See the training script for the complete FireDevourerEmbedder class definition
```

### Computing Embeddings

```python
def mean_pooling(model_output, attention_mask):
    """Apply mean pooling to get sentence embeddings."""
    token_embeddings = model_output[0]
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)

def get_embedding(text, model, tokenizer):
    """Get sentence embedding for a single text."""
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
    with torch.no_grad():
        outputs = model(**inputs)
    return mean_pooling(outputs, inputs["attention_mask"])

# Example
embedding = get_embedding("This is a sample sentence.", model, tokenizer)
print(f"Embedding shape: {embedding.shape}")  # [1, 384]
```

### Computing Similarity

```python
from torch.nn.functional import cosine_similarity

sentence1 = "A man is eating pizza"
sentence2 = "A person is eating food"

emb1 = get_embedding(sentence1, model, tokenizer)
emb2 = get_embedding(sentence2, model, tokenizer)

similarity = cosine_similarity(emb1, emb2)
print(f"Similarity: {similarity.item():.4f}")  # ~0.9448
```

### Task-Specific Predictions

```python
# After loading the full model with task heads:

def predict_nli(sentence1, sentence2, model, tokenizer):
    """Predict entailment relationship."""
    # Get embeddings for both sentences
    emb1 = get_embedding(sentence1, model, tokenizer)
    emb2 = get_embedding(sentence2, model, tokenizer)

    # Combine embeddings (concatenate with element-wise difference and product)
    combined = torch.cat([emb1, emb2, torch.abs(emb1 - emb2), emb1 * emb2], dim=-1)

    # Pass through NLI head
    logits = model.nli_head(combined)
    prediction = torch.argmax(logits, dim=-1)

    labels = ["Contradiction", "Neutral", "Entailment"]
    return labels[prediction.item()]

# Example
result = predict_nli("It's raining outside", "The weather is sunny", model, tokenizer)
print(f"NLI Prediction: {result}")  # Contradiction
```

## Evaluation Results

### Test Suite Statistics (20 diverse test cases)

**Cosine Similarity:**
| Statistic | Value |
|-----------|-------|
| Mean | 0.8001 |
| Std | 0.1562 |
| Min | 0.3139 |
| Max | 0.9831 |
| Median | 0.8149 |

**STS Score:**
| Statistic | Value |
|-----------|-------|
| Mean | 0.5672 |
| Std | 0.2270 |
| Min | 0.0182 |
| Max | 0.9468 |
| Median | 0.5788 |

### Example Predictions

| Sentence 1 | Sentence 2 | Cosine Sim | NLI | Domain |
|------------|------------|------------|-----|--------|
| "A man is eating pizza" | "A person is eating food" | 0.9448 | Entailment | General |
| "It's raining outside" | "The weather is sunny" | 0.7124 | Contradiction | Entailment |
| "How do I learn Python?" | "What's the best way to learn Python?" | 0.8915 | Entailment | Questions |
| "The quick brown fox jumps..." | "A fast brown fox leaps..." | 0.7837 | Entailment | General |

## Intended Use

### Best Use Cases
- **Semantic Search**: Finding similar documents or passages
- **Duplicate Detection**: Identifying duplicate questions or content
- **Paraphrase Mining**: Finding paraphrased text pairs
- **Clustering**: Grouping similar sentences or documents
- **Natural Language Inference**: Determining textual entailment

### Limitations
- **STS-B Performance**: The model shows lower performance on fine-grained semantic similarity regression (0.3366). For tasks requiring precise similarity scores, consider using dedicated STS models.
- **English Only**: Trained exclusively on English data.
- **Max Length**: 512 tokens maximum input length.
- **Adversarial Robustness**: While trained on PAWS adversarial data, performance on novel adversarial examples may vary.

## Training Loss Progression

| Epoch | STS Loss | NLI Loss | QQP Loss | PAWS Loss | MRPC Loss | Domain Loss | Total Loss |
|-------|----------|----------|----------|-----------|-----------|-------------|------------|
| 1 | 0.0073 | 0.2508 | 0.0742 | 0.0966 | 0.0287 | 0.0529 | 0.4977 |
| 2 | 0.0038 | 0.1970 | 0.0430 | 0.0638 | 0.0025 | 0.0196 | 0.3211 |
| 3 | 0.0031 | 0.1822 | 0.0221 | 0.0479 | 0.0009 | 0.0141 | 0.2631 |

## Citation

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

```bibtex
@misc{firedevourerembedder2025,
  author = {Asad, Zain},
  title = {FireDevourerEmbedder: Multi-Task Sentence Embeddings with RL-Adaptive Task Weighting},
  year = {2025},
  publisher = {Hugging Face},
  url = {https://huggingface.co/zainasad/FireDevourerEmbedder-RL-v3.6}
}
```

## Author

**Zain Asad**

## License

MIT License