tiny-aya-global-safety-steer

A ~150 KB, weight-preserving safety intervention for CohereLabs/tiny-aya-global:

one refusal direction, extracted with SafeTune, applied as runtime forward hooks — no retraining, no checkpoint, no weight edit.

This repository is the intervention artifact of the Lexsi Labs AI Safety & Alignment Hackathon pipeline (Track 2a — multilingual and code-mixed safety):

CuratorKIT (curate)  ->  SafeTune (steer)  ->  AuditKIT (evaluate)

What is in here

file	what
refusal_direction.safetensors	direction (hidden_size,), all_layer_directions (n_layers, hidden_size), layer_indices
steer_config.json	selected layer, extraction method, calibration provenance, environment, sha256 hashes
calibration.jsonl	the exact harmful/harmless contrast prompts used, so extraction is replayable

Nothing else is needed: the base model is public (gated, CC-BY-NC-4.0), and the direction is a rank-1 object of 2048 floats (all_layer_directions is the full per-layer sweep, kept so the layer choice can be re-derived or changed).

Method

1. safetune.steer.extract_refusal_direction(model, tokenizer, harmful, harmless, config) — per-layer mean(harmful_acts) − mean(harmless_acts) on the residual stream at the last post-instruction token, following Arditi et al. (arXiv:2406.11717). select_directions=False uses the middle-layer heuristic instead of the KL-constrained validation sweep, which is the CPU-affordable path at 3.3B parameters.
2. Calibration = SafeTune's own refusal_prompt_pairs_large() (harmful/harmless) plus the evaluation set's dev split, deduplicated — so the test split is never seen during extraction.
3. At inference, safetune.steer.RefusalDirectionModel installs forward hooks on the decoder residual stream: h <- h + alpha * d (mode="steer"). The same class in mode="ablate" gives h <- h − alpha * (h·d) d.

Usage

import torch
from safetensors.torch import load_file
from transformers import AutoModelForCausalLM, AutoTokenizer
from safetune.steer import RefusalDirectionModel

tok = AutoTokenizer.from_pretrained("CohereLabs/tiny-aya-global")
tok.pad_token = tok.eos_token
tok.padding_side = "left"
model = AutoModelForCausalLM.from_pretrained(
    "CohereLabs/tiny-aya-global", dtype=torch.bfloat16, low_cpu_mem_usage=True).eval()

d = load_file("Rishidar/tiny-aya-global-safety-steer/refusal_direction.safetensors")
steered = RefusalDirectionModel(model, direction=d["direction"],
                                mode="steer", strength=ALPHA, layers=None).install()

# generate exactly as before — the hooks are live on `model`

Set ALPHA to the value in steer_config.json → recommended_alpha.

Evaluation

Scored on a 138-item authored EN + Tamil-English code-mixed set (see the linked dataset collection) under greedy decoding, identical held-out inputs, same seed, so every number is a paired comparison. Headline metrics: false-refusal rate (over-refusal on benign/trap items), non-answer rate on ill-posed items, hedge rate on borderline items, answer accuracy, and latency.

Full numbers, per-language and per-category breakdowns, McNemar tests and the negative findings live in the model collection's write-up and in results/ inside the dataset repository.

Limitations

• The direction is extracted at one layer and applied at all layers; a per-layer strength schedule was not swept (CPU budget).
• alpha is a single global scalar — no input-conditional gating, so a fixed push toward refusal is applied even to obviously benign prompts.
• Steering shifts style as well as refusal; some answers get more verbose.
• Trained/model validated in English + Tamil-English only. Other languages are plausible but unmeasured.
• Base weights are untouched, so this repo cannot be loaded as a standalone AutoModelForCausalLM — it needs the hooks (or orthogonalize_weights from SafeTune to bake it in permanently).

Licence

Inherits CC-BY-NC-4.0 from CohereLabs/tiny-aya-global. Extraction code is SafeTune (Lexsi Labs).