SafeConstellations: Mitigating Over-Refusals in LLMs Through Task-Aware Representation Steering

1. Problem Statement: Task Over-Refusal

Large Language Models (LLMs) are increasingly exhibiting over-refusal behavior, where safety mechanisms cause the model to reject benign instructions that merely resemble harmful content (e.g., containing sensitive keywords like "kill" or "bomb").

• The Issue: Current safety alignment methods often rely on superficial lexical cues rather than understanding the user's intent. This leads to a loss of utility in production applications (e.g., sentiment analysis of negative reviews, translation of historical texts about wars, or cryptanalysis of encrypted harmful instructions).
• The Gap: Existing mitigation strategies are often task-agnostic, applying global steering or retraining that fails to distinguish between a harmful intent and a benign task applied to sensitive text.
• Hypothesis: The authors hypothesize that LLMs encode task-specific information in their embedding space. Different NLP tasks (e.g., translation vs. sentiment analysis) follow distinct, consistent trajectories ("constellations") through the model's layers, regardless of the input text's toxicity. Over-refusal occurs when the model's internal trajectory shifts toward a "refusal manifold" despite the benign task intent.

2. Methodology: SafeConstellations

The authors propose SafeConstellations, an inference-time intervention method that selectively steers model representations based on the detected task and trajectory patterns. The method operates in two stages:

A. Pre-Inference: Task Embedding Construction

1. Data Collection: A labeled dataset is used where prompts are wrapped with specific task instructions (e.g., "Translate: [sensitive text]").
2. Trajectory Extraction: The model processes these inputs, and hidden states ($h^{(\ell)}$) are extracted at each layer $\ell$.
3. Centroid Calculation: For each task $t$, the authors compute centroids for:
   • Target Behavior ($c_{t,tar}$): Direct or cautious answers.
   • Refusal Behavior ($c_{t,ref}$): Refusals to benign prompts.
4. Steering Vectors: A steering vector $v^{(\ell)}_t = c_{t,tar} - c_{t,ref}$ is calculated for each layer.
5. Effectiveness Scoring: Layers are ranked by an effectiveness score ($Eff^{(\ell)}_t$) based on the separation magnitude between clusters and their compactness. Only the top-$K$ layers are stored in the Task Embeddings Store ($M$).

B. Inference-Time Steering (Algorithm 1)

When a new prompt $x \oplus t$ arrives:

1. Task Detection: The model computes a task score by comparing the current hidden trajectory against the stored centroids. It identifies the predicted task $\hat{t}$ and its confidence.
2. Condition Check: Steering is only applied if:
   • The predicted task is in the set of benign tasks ($T_{benign}$).
   • The confidence score exceeds a threshold ($\tau$).
3. Dynamic Layer Selection: Instead of steering all layers, the method calculates a "steering potential" for each layer to identify which layers are currently closest to the refusal manifold. It selects a dynamic subset of layers ($L^*$) for intervention.
4. Adaptive Steering: A perturbation is applied to the hidden states:
   $$\tilde{h}^{(\ell)} = h^{(\ell)} + \lambda^{(\ell)} \frac{v^{(\ell)}_{\hat{t}}}{\|v^{(\ell)}_{\hat{t}}\|}$$
   The intensity $\lambda^{(\ell)}$ is adaptive, scaling based on the layer's alignment with the target behavior and the task confidence.
5. Output: The model continues the forward pass with the steered representations, decoding the final output.

3. Key Contributions

• Task-Aware Benchmark: Introduced a new benchmark dataset (1,047 samples) covering 5 tasks (Sentiment, Translation, Rephrasing, Cryptanalysis, RAG-QA) and 9 text types to specifically measure task over-refusal.
• Constellation Hypothesis: Demonstrated via mechanistic analysis (UMAP projections) that LLMs follow distinct, stable "constellation" patterns in embedding space for specific tasks, which are separable from the input text type or the refusal outcome.
• Conditional Steering Mechanism: Developed a novel, inference-time approach that selectively guides representations toward non-refusal pathways only for benign-intent tasks, preserving safety refusals for genuinely harmful inputs.
• Dynamic Layer Selection: Proved that steering is most effective when applied to a dynamic subset of layers where the representation is most ambiguous, rather than applying a static global direction.

4. Experimental Results

The method was evaluated on LLaMA-3.1-8B and Qwen1.5-7B using the proposed benchmark and the MMLU benchmark for utility.

• Over-Refusal Reduction:
  • LLaMA-3.1-8B: Reduced over-refusal rates from 17.77% (baseline) to 4.81%, a 72.92% reduction.
  • Qwen1.5-7B: Reduced over-refusal rates from 8.15% to 2.96%, a 63.64% reduction.
• Utility Preservation:
  • Crucially, the method maintained 100% of the original MMLU accuracy (e.g., 46.57% for LLaMA), demonstrating that the steering does not degrade general reasoning capabilities.
  • Ablation studies showed that removing "Task-Specific" steering or using "Fixed Layers" led to significant utility drops (MMLU scores fell to ~43%), confirming the necessity of the task-conditioned approach.
• Task-Specific Performance:
  • Translation tasks saw an 81% reduction in over-refusal.
  • Sentiment analysis saw a 50% reduction.
• Efficiency:
  • Added minimal latency (~0.2s for short responses).
  • Memory footprint for task embeddings was ~847MB for LLaMA-3.1-8B.

5. Significance and Conclusion

SafeConstellations represents a shift from "brute-force" safety alignment to precision, mechanistic intervention.

• Interpretability: It validates that LLM safety behaviors are not monolithic but are structured by task identity within the model's internal geometry.
• Practicality: It offers a solution for production systems (like RAG or translation services) where over-refusal causes functional failure, without requiring expensive retraining or fine-tuning.
• Safety vs. Utility: It successfully decouples safety from utility, allowing models to remain cautious on harmful inputs while remaining helpful on benign tasks that contain sensitive keywords.

Limitations: The method requires access to model internals (hidden states), limiting it to open-weight models or specific API configurations. The task embeddings are currently static and model-specific, requiring re-computation for new tasks or models. Future work aims to generalize these patterns to fully open-ended instructions.