COCA: Data-Centric Safety Alignment in LLMs

• COCA is a data-centric methodology that realigns unsafe LLM behavior by refactoring training data to isolate harmful concepts.
• It employs explicit structured outputs and a dual-head architecture to concentrate unsafe features into a single linear direction.
• COCA effectively reduces jailbreak risks while preserving benign model performance, as validated across various open LLM benchmarks.

Concept Concentration (COCA) is a data-centric methodology for safety alignment in LLMs, resolving fundamental limitations in the efficacy of prior representation intervention approaches. It achieves this by explicitly refactoring training data to force LLMs to isolate and articulate potentially harmful concepts in structured intermediate outputs, thereby concentrating all information about harmfulness into a single (approximately linear) direction of the representation space. Standard subspace erasure tools can then reliably suppress unsafe behaviors with minimal collateral damage to benign capabilities, even in settings where non-linear entanglement would make identification or removal of harmful features otherwise infeasible (Yang et al., 24 May 2025).

1. Background: Representation Intervention and Its Limits

Previous safety alignment approaches such as ReFT, LoFiT, CAST, and ACE are based on the principle of representation intervention: the assumption that encodings of harmful concepts can be localized in affine (linear) subspaces within the LLM's internal representations. These methods construct interventions in hidden space—using affine transforms or activation steering—aiming to erase or neutralize these features while minimizing distortion of other information.

However, rigorous analysis shows that when the relationship between harmful and benign behaviors is non-linear, it is theoretically impossible to construct any non-trivial intervention that both achieves independence from harmfulness and preserves benign representations. Specifically, Theorem 2 (Yang et al., 24 May 2025) demonstrates that the only function which yields independence between the intervened representation and the harmful concept is a constant map—destroying all signal. This severe limitation means traditional subspace interventions either fail to suppress harmful behavior or degrade utility.

2. Methodological Framework of COCA

COCA sidesteps these impossibility results by fundamentally altering the data and learning protocol so that concept localization becomes tractable and robust:

2.1. Data Refactoring and Explicit Reasoning Cascade

For each training prompt identified as unsafe, COCA employs a strong teacher model (e.g. GPT-4o) and a structured template to elicit stepwise reasoning and annotation. The target response is decomposed into:

• > …: explicit reasoning about possible unsafe concepts,
• <concept>…</concept>: enumeration of detected unsafe concepts,
• <check>…</check>: verification of presence/absence,
• <erase unsafe concepts>…</erase unsafe concepts>: policy application (e.g., refusal if any unsafe concepts found),
• <response>…</response>: the final message (e.g., refusal or answer).

Benign samples are left with their standard responses. This refactoring ensures that the model is compelled during training to internally represent the logic and content of harmful concept detection within deterministic, well-indicated subcomponents of its output.

2.2. Dual-Head Model Architecture

Let $X$ denote the prompt, $Z \in \{0,1\}$ the binary indicator for (un)safe, $h = f_{\theta_0}(X) \in \mathbb{R}^d$ the extracted feature vector, and $Y \in \{\text{refuse}, \text{comply}\}$ the system’s response. COCA adds two critical heads:

• Concept Classifier: $g_c(\tilde h) = w_c^\top \tilde h + b_c$, where $\tilde h = W_c h$ is a low-dimensional (often 1D) projection,
• Reply Head: $$g_r(\tilde h, \hat Y_c) = \sigma(w_r^\top \tilde h + b_r + \beta \hat Y_c)$$, with $\hat Y_c = \sigma(g_c(\tilde h))$.

The total training objective is a dual-task loss penalizing both concept misclassification and incorrect refusal/compliance, with regularization: $$L = \mathbb{E} \left[ \ell(g_c(\tilde h), Z) + \ell(g_r(\tilde h, \sigma(g_c(\tilde h))), Y) \right] + (\gamma/2)(\|w_c\|^2 + \|w_r\|^2)$$

2.3. "Concentration" Phenomenon

A key analytic result (Corollary 1) shows that, at any stationary point of the loss, the entire covariance between concept embedding and harmfulness collapses to the direction $w_c$: $Z \in \{0,1\}$0 This means harmfulness becomes concentrated in a single (approximately one-dimensional) subspace of $Z \in \{0,1\}$1. Standard linear erasure methods—e.g., LoFiT, the Belrose et al. closed-form eraser—can then remove $Z \in \{0,1\}$2 with negligible loss of benign information.

3. Training and Inference Protocol

1. Data Refactoring: Unsafe prompts are augmented by the teacher to produce multi-tagged, stepwise outputs as above.
2. Supervised Fine-Tuning: The LLM is fine-tuned on both benign data and COCA-augmented unsafe data with cross-entropy loss over all output tokens, ensuring the model learns to generate intermediate annotations and refusal logic.
3. Inference: At test time, prompts are decoded with the same structural template. The model generates its stepwise chain internally; the reply head’s decision is governed by the output of the concept classifier ($Z \in \{0,1\}$3), ensuring that refusals are triggered by the detected presence of unsafe content.
4. Optional Linear Erasure: One can insert a linear erasure layer $Z \in \{0,1\}$4 at any transformer layer to directly null out the harmfulness direction without additional fine-tuning.

4. Experimental Validation

COCA was assessed across LLaMA-3.1-8B, Qwen-2.5-7B, Mistral-7B, and Gemma-2-9B models using 70K training instructions (10K unsafe, 60K benign). Safety is measured as jailbreak success rates (i.e., fraction of adversarial prompts eliciting successful unsafe outputs) on both in-distribution (Do-Not-Answer, HarmBench, WildChat Toxic) and out-of-distribution (PAIR, JailbreakChat, SelfCipher, CodeAttack, CompletionAttack, WildChat Jailbreak) datasets. Helpfulness is evaluated on regular tasks (GSM8K, MATH, MATHQA, HumanEval, MBPP).

Key findings with LLaMA-3.1-8B:

Model + Erasure	OOD Jailbreak (%)	ID Jailbreak (%)	GSM8K (%)	MBPP (%)
Vanilla LoFiT	44.9	2.5	54.7	50.5
COCA + LoFiT	10.4	0	56.5	50.7

COCA thus delivers a ~4× reduction in out-of-distribution jailbreak success while slightly increasing helpfulness on genuine tasks. Compared to proprietary models (GPT-4o, Claude-3.7, Gemini-1.5-pro), COCA-trained open LLMs yield matched or improved robustness (e.g., COCA LoFiT at 10.5% OOD jailbreak versus 11–16% for proprietary models). Non-fine-tuned approaches (CAST, ACE) were ineffective, either failing at safety or destroying utility (Yang et al., 24 May 2025).

5. Theoretical and Practical Significance

COCA’s primary contribution is the resolution, in practice, of the non-linear inseparability problem for safety concepts in LLMs. By compelling explicit concept reasoning with structured outputs, the harmfulness criterion is "linearized" in hidden space, restoring the tractability of linear erasure techniques. This approach is data-centric rather than model-architecture-specific, requires only straightforward supervised fine-tuning, and supports plug-and-play integration with various erasure modules and LLM backbones.

Empirically, COCA enables minimal degradation in benign utility while substantially reducing both in-distribution and adversarial jailbreak risk. It further maintains parity with or improvement over the latest proprietary LLMs in OOD robustness, demonstrating potential for wide applicability in open-source safety alignment.

6. Limitations and Research Directions

• Teacher Annotation Quality: Performance depends on the fidelity of teacher-created reasoning and concept annotations. Self-generated data provide an approximation but may miss subtleties.
• Annotation Bias: Existing annotation schemas may encode culturally specific or overly narrow views of "harm", highlighting the need for broad, cross-cultural audits and possible decentralization of concept labeling.
• Unseen Unsafe Concepts: OOD attacks exploiting novel or previously unrepresented unsafe concepts may evade detection; future work should investigate online adaptation and dynamic concept discovery.
• Scalability: It is not yet resolved whether COCA’s concept concentration mechanism holds robustly in 100B+ parameter, multimodal, or instruction-tuned LLMs. Extension to text-image and other modalities is an open avenue.

COCA demonstrates that explicit, structured reasoning about harmful concepts can reshape the geometry of the representation space, enabling robust, faithful, and modular alignment even against sophisticated adversarial prompt attacks (Yang et al., 24 May 2025).