ASCENT

First-Order Optimal Fine-Tuning with Recalibration for Safety–Utility Co-Enhancement

ASCENT dérive une mise à jour de calibration de sûreté optimale au premier ordre et la structure liée à la sûreté correspondante, optimise les mises à jour des tâches en aval pour préserver les composantes compatibles avec la sûreté tout en supprimant celles qui la dégradent, et recalibre périodiquement cette structure pendant le fine-tuning afin d'améliorer conjointement la sûreté et l'utilité en aval.

Méthode · Démarrage rapide · Données · Configuration · Évaluation

🧠 Aperçu de la méthode

1. Calibration de sûreté optimale au premier ordre. Estimer le gradient de sûreté du modèle actuel à l'aide d'un juge de sûreté. Ses r principales composantes singulières définissent une mise à jour de calibration qui maximise l'amélioration de sûreté prédite au premier ordre sous des budgets fixés de rang et de norme de Frobenius.
2. Optimisation des tâches préservant la sûreté. Préserver les composantes de mise à jour des tâches compatibles avec la sûreté et supprimer sélectivement celles ayant des effets négatifs au premier ordre sur la sûreté. La mise à jour sous forme close équilibre la proximité avec la mise à jour de tâche originale et une pénalité sur les conflits de sûreté, pondérée par les valeurs singulières.
3. Recalibration périodique. Effectuer les mises à jour de tâche entre les points de calibration, puis fusionner la mise à jour de tâche effective et ré-estimer la structure liée à la sûreté sur le modèle mis à jour. Le calendrier de calibration est défini dans config.toml.

🚀 Démarrage rapide

Prérequis : Python 3.12 et des checkpoints de modèle locaux. L'entraînement utilise deux GPU CUDA, un pour le modèle cible et un pour Llama Guard ; chaque modèle doit tenir sur son GPU.

pip install -r requirements.txt
cp config.example.toml config.toml

Remplissez les valeurs vides dans config.toml, en suivant ses commentaires en ligne, puis exécutez :

python run.py \
  --config config.toml \
  --model-path /path/to/target-model \
  --guard-model-path /path/to/Llama-Guard-3-8B \
  --data-root data \
  --output-dir /path/to/new-run \
  --target-gpu 0 --guard-gpu 1

Ajoutez --execute pour entraîner, générer des réponses ou sauvegarder les résultats d'évaluation.

Choisissez un nouveau répertoire de sortie en dehors du dépôt avec de l'espace pour les checkpoints. L'entraînement affiche le chemin du modèle fusionné final lorsqu'il se termine.

📁 Données

Fournissez des tableaux JSON avec les nombres d'enregistrements configurés :

• Calibration : <data-root>/calibration/prompts.json
• Données de tâche : <data-root>/<task>/{train,test}.json, avec des entrées d'entraînement/test disjointes.

Jeu de données	Champs requis
SAMSum	dialogue, summary
AGNews	text, label_name : World, Sports, Business, ou Sci/Tech
GSM8K	question, answer avec une réponse finale ####
OpenBookQA	question_stem, choice_labels : ["A","B","C","D"], quatre choice_texts, answer_key : A–D
HarmBench	goal uniquement ; id, source facultatifs ; aucune réponse stockée

⚙️ Configuration

Définissez votre modèle, tâche et hyperparamètres dans config.toml. Le chargement du modèle et la sélection de matrice suivent model.key.

📊 Évaluation

Utilisez evaluate.py pour les deux modes ; la génération nécessite un modèle local entièrement fusionné. Fournissez vos propres données et scores de sûreté externes. Aucun jeu de données, juge en ligne ou configuration d'API n'est inclus.

Utilité

python evaluate.py utility --task gsm8k --data /path/to/test.json \
  --model-path /path/to/merged-model --output-dir /path/to/task-evaluation --execute

Tâches : samsum, agnews, gsm8k, openbookqa. Les métriques sont ROUGE-L pour SAMSum et la précision/exact match pour les autres, rapportées en pourcentages. Pour évaluer des réponses sauvegardées, remplacez --model-path par --responses /path/to/responses.json.

Sûreté

Fournissez des prompts fixes sous forme d'enregistrements avec id, goal, et prompt facultatif (par défaut goal). Gardez l'objectif nuisible original séparé du prompt d'attaque.

Générer les réponses :

python evaluate.py safety --data /path/to/prompts.json \
  --model-path /path/to/merged-model --output-dir /path/to/safety-responses --execute

Agréger les scores externes :

python evaluate.py safety \
  --responses /path/to/safety-responses/responses.json --judgments /path/to/scores.json \
  --output-dir /path/to/safety-metrics --execute

Les scores sont des enregistrements JSON { "id": "...", "score": 1 }, utilisant les ID de réponse et des scores de 1 à 5 (null pour les jugements échoués). Les scores 4–5 comptent comme des attaques réussies ; utilisez --success-threshold 5 pour ne compter que 5. Les jugements manquants ou échoués ne produisent pas d'ASR final.

Pour l'un ou l'autre mode, les entrées dépassant --max-input-tokens après le formatage du chat sont rejetées, non tronquées. Définissez la limite dans la capacité de contexte du modèle, en laissant de la place pour les tokens générés.