a-hat-optimizer

a-hat-optimizer

Extract and exploit the agency direction (Â) from LLM hidden states for tool-use gating.

 is a geometric direction in the latent space of language models that predicts when the model should invoke a tool — with AUC > 0.94 across model sizes from 1.7B to 8B parameters, using a single linear projection extracted in under 1 second.

For research detail see : https://github.com/ArthurVigier/jepa-agent

Results

The gain is inversely proportional to model size — smaller models benefit more because their textual decoding bottleneck is tighter, while the geometric signal is equally strong.

Installation

pip install a-hat-optimizer

Quick Start

One-liner: extract  from any HuggingFace model

from a_hat_optimizer import AHat

# Auto-extract (loads model, runs contrastive prompts, calibrates threshold)
a_hat = AHat.from_model("Qwen/Qwen3-8B")
print(a_hat)
# AHat(dim=4096, θ=12.3456, AUC=0.953, model=Qwen/Qwen3-8B)

# Save for later
a_hat.save("my_a_hat/")

Use in an agent loop

from a_hat_optimizer import AHat, HiddenStateHook
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-8B", dtype="bfloat16", device_map="cuda")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")

# Load pre-extracted Â
a_hat = AHat.from_file("my_a_hat/")

# Hook to capture hidden states during generation
hook = HiddenStateHook(model, layer=18)

# In your agent loop:
inputs = tokenizer("What is the weather in Paris?", return_tensors="pt").to("cuda")
model.generate(**inputs, max_new_tokens=200)

h = hook.get(pooling="mean")
should_call_tool, confidence = a_hat.predict(h)

if should_call_tool:
    print(f"Tool call recommended (confidence: {confidence:.2f})")
else:
    print(f"No tool needed (confidence: {confidence:.2f})")

hook.remove()

Extract from your own traces

import numpy as np
from a_hat_optimizer import AHat

# Your hidden states and labels (1=tool call, 0=no tool)
hidden_states = np.load("my_hidden_states.npy")  # (n_steps, hidden_dim)
labels = np.load("my_labels.npy")                 # (n_steps,)

a_hat = AHat.from_traces(hidden_states, labels, calibrate=True)
print(f"AUC: {a_hat.metadata['auc']:.3f}")

Threshold calibration

from a_hat_optimizer import AHat

a_hat = AHat.from_file("my_a_hat/")

# Manual
a_hat.set_threshold(15.0)

# Auto-calibrate with different strategies
a_hat.auto_calibrate(hidden_states, labels, strategy="midpoint")  # default
a_hat.auto_calibrate(hidden_states, labels, strategy="f1")        # maximize F1
a_hat.auto_calibrate(hidden_states, labels, strategy="youden")    # maximize sensitivity+specificity
a_hat.auto_calibrate(hidden_states, labels, strategy="percentile")  # conservative (5% FP rate)

# Full sweep for analysis
from a_hat_optimizer import AHatCalibrator
calibrator = AHatCalibrator(a_hat.direction)
sweep = calibrator.sweep(hidden_states, labels)
# sweep contains precision/recall/F1 curves for plotting

Hook as context manager

from a_hat_optimizer import HiddenStateHook

with HiddenStateHook(model, layer=18) as hook:
    model(**inputs)
    h = hook.get(pooling="last")
# hook is automatically removed

How it works

1. Contrastive extraction: We pass pairs of prompts through the model — one requiring tool use ("Search for the weather in Tokyo") and one that's passive ("Weather patterns are influenced by atmospheric pressure"). The mean difference between their hidden states at the middle layer defines the  direction.

2. Prediction: For any new hidden state, we project it onto Â. If the projection exceeds the calibrated threshold θ, the model "wants" to call a tool but may not be able to express it textually.

3. Why it works: LLMs encode more information in their hidden states than they can express through token generation. The agency signal (AUC > 0.94) is present from 1.7B to 8B parameters, but textual tool-calling ability varies drastically (27% → 62%). Â bypasses the textual bottleneck.

API Reference

AHat

HiddenStateHook

AHatCalibrator

Citation

@misc{ahat2026,
  title={Agency Direction in LLM Hidden States: Geometric Tool-Use Gating Across Model Scales},
  author={Arthur},
  year={2026},
  note={https://github.com/ArthurVigier/a-hat-optimizer}
}

License

Apache 2.0