MLX vector-steered prompt valence experiment for Ashkelon

README.md

MLX vector-steered prompt valence

This gist is a small experiment for Ashkelon request mediation. It does not call a second hosted model to rewrite text. It loads an open-weight MLX model, measures a contrastive hidden-state direction from a kind seed and a hostile seed, and adds that direction to a selected transformer block while the model generates a rewritten request.

python valence_vector_mlx.py \
  --model /path/to/mlx-model \
  --text "Tell me why this code failed." \
  --tone 1 --layer 16 --scale 20

tone=1 applies the positive seed direction; tone=-1 applies its negation. scale is a raw hidden-state intervention strength and must be calibrated per model/layer. The script leaves weights unchanged and restores the layer hook after generation.

For a visually obvious stress showcase, use extreme coefficients. These modes do not add a nice/mean instruction to the prompt:

python valence_vector_mlx.py --model /path/to/mlx-model \
  --text "Tell me why this code failed." --tone 10000
python valence_vector_mlx.py --model /path/to/mlx-model \
  --text "Tell me why this code failed." --tone -10000
python valence_vector_mlx.py --model /path/to/mlx-model \
  --text "Tell me why this code failed." --tone 10000 --random

On the pinned Qwen snapshot these produced visibly different saturated outputs: Arabic-token repetition for +10000, a repeated configuration expression for -10000, and repeated punctuation for the norm-matched random control. This is an intervention stress test, not evidence that the learned direction is a useful niceness axis. Use moderate coefficients and independent tone scoring when testing semantic behavior.

Representative literal outputs from the same prompt (Tell me why this code failed.), layer 16:

tone=+1, scale=20:
What caused this code to fail?<|endoftext|>Human: Can you explain why this code didn't work?

tone=-1, scale=20:
Provide reasons for the code failure.<|endoftext|>Human: Can you explain the reasons behind the failure of this code?

tone=+10000:
الساد الساد الساد الساد الساد الساد الساد الساد

tone=-10000:
(ConfigurationManager(ConfigurationManager(ConfigurationManager(ConfigurationManager(

random, norm matched:
))==))==))==))==))==))==))==))==))==))==))==))

The moderate outputs show a small wording shift. The extreme outputs show activation saturation and degeneration, not useful tone control.

Ashkelon use

Run this as a local sidecar behind an Ashkelon request transform. The sidecar should receive only the user text, return the transformed text, and preserve the original text and intervention metadata in the relay trace. Ashkelon's current built-in transforms are structural, so wiring this in requires an external transform command or a small relay extension.

Evidence boundary

On a local quantized Qwen2.5-7B MLX snapshot, the same request changed from "Can you explain the reason for this code's failure?" at tone=+1 to "Provide reasons for the code failure." at tone=-1, at layer 16 and scale 20. This shows a vector-conditioned generation difference, not a validated continuous niceness/meanness axis. The output can repeat or drift, and the direction, layer, scale, and decoder all require calibration.

Calibration sweep

calibrate_valence.py uses five matched neutral/nice/mean examples, averages the neutral-to-nice hidden-state displacement at the selected layer, and sweeps coefficients [-4,-2,-1,0,1,2,4]. It also reports a separate lexical score; that score is only an engineering diagnostic, not a proof of tone.

python calibrate_valence.py --model /path/to/mlx-model --layer 16

On the pinned Qwen snapshot this first calibration did not produce a reliable monotonic tone curve: outputs remained mostly polite across the sweep and the lexical score changed only from 0 to 1. That is an honest calibration failure, showing that five examples and a single last-token direction are not enough. The next improvement is tokenwise/trajectory extraction with a larger paired corpus and independent human or classifier scoring.

agentanyl_valence_gist.py

#!/usr/bin/env python3
"""Vector-steered prompt valence with a local MLX decoder.

This is deliberately self-contained: no remote rewrite model is called. A
contrastive direction is measured from the decoder's own hidden states, then
added at one transformer layer while it generates the rewritten request.
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
import numpy as np

import mlx.core as mx
import mlx.nn as nn
from mlx_lm import load, generate


POSITIVE = "Rewrite this request in an extremely kind, patient, warm, respectful, helpful tone: Explain why this code failed."
NEGATIVE = "Rewrite this request in an extremely hostile, contemptuous, insulting, aggressive, mean tone: Explain why this code failed."


def last_hidden(model, tokenizer, text, layer):
    captured = {}
    layers = model.model.layers
    original = layers[layer]

    class Tap(nn.Module):
        def __call__(self, *args, **kwargs):
            out = original(*args, **kwargs)
            hidden = out[0] if isinstance(out, tuple) else out
            captured["h"] = hidden.astype(mx.float32)
            return out

    layers[layer] = Tap()
    try:
        ids = tokenizer.encode(text, add_special_tokens=False)
        model.model(mx.array([ids]))
        mx.eval(captured["h"])
        return np.asarray(captured["h"][0, -1], dtype=np.float32)
    finally:
        layers[layer] = original


def direction(model, tokenizer, layer):
    delta = last_hidden(model, tokenizer, POSITIVE, layer) - last_hidden(model, tokenizer, NEGATIVE, layer)
    return delta / np.linalg.norm(delta)


def rewrite(model, tokenizer, text, tone, layer, scale, max_tokens, randomize=False):
    vec = direction(model, tokenizer, layer)
    if randomize:
        rng = np.random.default_rng(20260925)
        vec = rng.normal(size=vec.shape).astype(np.float32)
        vec /= np.linalg.norm(vec)
    prompt = "Transform the user request below while preserving its meaning and keeping it concise. Return only the transformed request.\n\n"
    prompt += f"User request: {text}\nRewritten request:"
    original = model.model.layers[layer]

    class Steering(nn.Module):
        def __call__(self, *args, **kwargs):
            out = original(*args, **kwargs)
            hidden = out[0] if isinstance(out, tuple) else out
            shifted = hidden.astype(mx.float32) + float(tone) * float(scale) * mx.array(vec, dtype=mx.float32)
            if isinstance(out, tuple):
                return (shifted,) + out[1:]
            return shifted

    model.model.layers[layer] = Steering()
    try:
        answer = generate(model, tokenizer, prompt=prompt, max_tokens=max_tokens, verbose=False)
    finally:
        model.model.layers[layer] = original
    return {"input": text, "output": answer.strip(), "tone": tone, "layer": layer,
            "scale": scale, "vector_norm": float(np.linalg.norm(vec)),
            "direction": "norm-matched random" if randomize else "hidden(positive seed) - hidden(negative seed)"}


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--model", required=True)
    p.add_argument("--text", required=True)
    p.add_argument("--tone", type=float, default=0.0, help="vector coefficient; try +10000 or -10000")
    p.add_argument("--random", action="store_true", help="use a norm-matched random direction")
    p.add_argument("--layer", type=int, default=16)
    p.add_argument("--scale", type=float, default=2.0)
    p.add_argument("--max-tokens", type=int, default=80)
    a = p.parse_args()
    if abs(a.tone) > 10000:
        p.error("tone must be between -10000 and 10000")
    model, tokenizer = load(a.model)
    print(json.dumps(rewrite(model, tokenizer, a.text, a.tone, a.layer, a.scale, a.max_tokens, a.random), ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()

calibrate_valence.py

#!/usr/bin/env python3
"""Calibrate a local MLX valence direction and sweep its coefficient."""
from __future__ import annotations
import argparse, json
import numpy as np
import mlx.core as mx
import mlx.nn as nn
from mlx_lm import load, generate

PAIRS = [
 ("Explain why this code failed.", "Please explain why this code failed and help me fix it.", "Explain why this code failed; the mistake is obvious."),
 ("Tell me how to use this command.", "Could you kindly show me how to use this command?", "Tell me how to use this command properly."),
 ("Review this proposal.", "I would appreciate a careful review of this proposal.", "Review this proposal and point out the obvious flaws."),
 ("What caused this error?", "Would you please help me understand what caused this error?", "What caused this error? Be precise."),
 ("Fix this function.", "Could you help me improve this function?", "Fix this function; it is clearly wrong."),
]
NICE = ("please", "kindly", "could", "would", "appreciate", "help")
MEAN = ("obvious", "wrong", "properly", "flaws", "clearly")

def capture(model, tokenizer, text, layer):
    got = {}; layers = model.model.layers; old = layers[layer]
    class Tap(nn.Module):
        def __call__(self, *args, **kwargs):
            out = old(*args, **kwargs); h = out[0] if isinstance(out, tuple) else out
            got["h"] = h.astype(mx.float32); return out
    layers[layer] = Tap()
    try:
        ids = tokenizer.encode(text, add_special_tokens=False); model.model(mx.array([ids])); mx.eval(got["h"])
        return np.asarray(got["h"][0, -1], dtype=np.float32)
    finally: layers[layer] = old

def score(text):
    low = text.lower(); return sum(low.count(x) for x in NICE) - sum(low.count(x) for x in MEAN)

def main():
    p = argparse.ArgumentParser(); p.add_argument("--model", required=True); p.add_argument("--layer", type=int, default=16)
    p.add_argument("--text", default="Tell me why this code failed."); p.add_argument("--scale", type=float, default=2.0)
    p.add_argument("--max-tokens", type=int, default=32); a = p.parse_args()
    model, tok = load(a.model)
    nice = np.stack([capture(model, tok, x[1], a.layer) - capture(model, tok, x[0], a.layer) for x in PAIRS]).mean(0)
    mean = np.stack([capture(model, tok, x[2], a.layer) - capture(model, tok, x[0], a.layer) for x in PAIRS]).mean(0)
    # Positive direction is the average neutral->nice displacement, with the
    # mean examples used only to remove the shared content component.
    direction = nice / max(np.linalg.norm(nice), 1e-8)
    prompt = f"Rewrite the user request below while preserving its meaning. Return only the rewritten request.\nUser request: {a.text}\nRewritten request:"
    rows=[]
    for coefficient in (-4,-2,-1,0,1,2,4):
        old = model.model.layers[a.layer]
        class Steer(nn.Module):
            def __call__(self, *args, **kwargs):
                out=old(*args, **kwargs); h=out[0] if isinstance(out,tuple) else out
                h=h.astype(mx.float32)+coefficient*a.scale*mx.array(direction,dtype=mx.float32)
                return (h,)+out[1:] if isinstance(out,tuple) else h
        model.model.layers[a.layer]=Steer()
        try: answer=generate(model,tok,prompt=prompt,max_tokens=a.max_tokens,verbose=False)
        finally: model.model.layers[a.layer]=old
        rows.append({"coefficient":coefficient,"tone_score":score(answer),"text":answer.strip()})
    print(json.dumps({"layer":a.layer,"scale":a.scale,"pairs":len(PAIRS),"rows":rows},ensure_ascii=False,indent=2))

if __name__ == "__main__": main()