Multi-Task Learning via Model Merging for Toxicity and Spam Detection

Sapienza University of Rome - Natural Language Processing Course (HWp) > 📄 Final Report: The comprehensive scientific discussion and analysis can be found in the Report - MTL Toxicity and Spam Detection.pdf file located in the root directory.

📖 Abstract

This project addresses a core challenge in Natural Language Processing: mitigating catastrophic forgetting without the computational cost of joint-training. We explore Task Arithmetic (Model Merging) as a parameter-efficient solution for Multi-Task Text Classification.

The goal is to merge the capabilities of a Transformer-based encoder (distilbert-base-uncased) to simultaneously resolve two semantically distant NLP tasks: Toxicity Detection and SMS Spam Detection. By fine-tuning separate models and fusing their extracted "Task Vectors" via linear interpolation, the project investigates how linguistic representations interact within the model's latent space. The merged architecture is evaluated against both the single-task encoders and a Zero-Shot Generative LLM baseline (facebook/bart-large-mnli). By tuning the interpolation hyperparameter ($\alpha$), the study quantifies parameter interference and finds the optimal representation trade-off when merging orthogonal linguistic domains.

📊 Datasets

We selected two well-established benchmarks representing different classification complexities:

1. TweetEval (Hate Subset): Used to evaluate Toxicity Detection, representing a complex, noisy, and highly subjective NLP task.
2. SMS Spam Collection: Used to evaluate Spam Detection, representing a more rigid, pattern-based classification task.

🧠 Methodology & Models

• Base Architecture: DistilBERT (Sanh et al., 2019), a distilled version of BERT.
• Models A & B (Task-Specific Encoders): Independent models obtained by fine-tuning the base architecture separately on the two datasets. These serve as the specialized upper-bound performance benchmarks.
• Merged Architecture: The single parameter-efficient model resulting from the linear combination of the two task vectors ($\tau_{Toxicity} + \tau_{Spam}$). To address parameter interference, we evaluate both Standard Model Merging (Ilharco et al., 2022) and the TIES-Merging technique (Yadav et al., 2023).
• Baseline System (Zero-Shot Generative LLM): An out-of-the-box bart-large-mnli model used in Zero-Shot mode (Yin et al., 2019) to compare specialized representations against broad parametric knowledge.

🏆 Key Findings

• Resolved Domain Collapse: While standard cross-task evaluation resulted in complete predictive failure, TIES-Merging ($\alpha=1.0$) successfully harmonized the orthogonal domains.
• Emergent Regularization: Injecting the rigid Spam task vector into the Toxicity model acted as a powerful regularizer, pushing the Toxicity Macro F1 to 0.59 (surpassing the specialized upper-bound of 0.48).
• LLM Outperformance: The lightweight merged DistilBERT outperformed a massive 400-million parameter Zero-Shot LLM (BART-Large) across both tasks.
• Hardware Efficiency: The unified multi-task encoder requires exactly 255.41 MB of VRAM, saving 50% in deployment memory costs compared to a traditional dual-model routing system.

📂 Repository Structure

The codebase is modular and automatically generates a structured analytics directory upon execution:

.
├── Report - MTL Toxicity and Spam Detection.pdf     # Final academic report
├── README.md                       
├── src/                            # Source Code
│   ├── dataset.py                  # Dataloaders and HuggingFace dataset processing
│   ├── model.py                    # Base architecture initialization
│   ├── train.py                    # Fine-tuning, evaluation, and latent space extraction
│   ├── merging.py                  # Task Arithmetic, TIES-Merging, and Cosine Similarities
│   ├── baseline_llm.py             # Zero-shot generative LLM evaluation
│   ├── plot_results.py             # High-resolution academic plotting suite
│   └── main.py                     # Orchestrator script
│
└── results/                        # Automatically generated output directory
    ├── 00_Base_Model/              # Untrained baseline metrics
    ├── 01_Task_A_Toxicity/         # Fine-tuned System A
    ├── 02_Task_B_Spam/             # Fine-tuned System B
    ├── 03_Cross_Task_Evaluation/   # Zero-shot cross-domain testing
    ├── 04_Merged_Standard/         # Results for Standard Task Arithmetic (alphas 0.0 to 1.0)
    ├── 05_Merged_TIES/             # Results for TIES-Merging (alphas 0.0 to 1.0)
    ├── 06_Baseline_LLM/            # Generative BART-Large evaluation
    ├── Plots/                      # 8 high-res charts (Trade-offs, ROCs, Heatmaps, T-SNE)
    └── [Raw Data Files]            # efficiency_report.txt, tsne_data.npz, .json analytics

(Note: Each subfolder in results/ contains detailed report.txt files, Confusion Matrices, ROC-AUC curves, and Error Analysis CSVs tracking specific misclassifications).

🚀 Execution & Reproducibility

Global random seeds (seed=42) are strictly enforced for 100% reproducibility.

1. Install dependencies:

pip install torch transformers datasets scikit-learn matplotlib seaborn

2. Run the full experimental pipeline:

python src/main.py

3. Regenerate plots (Optional): If raw data files exist in results/, you can regenerate the plot suite without re-training:

python src/plot_results.py

---

Project developed by Lorenzo Musso & Giulia Pietrangeli