optimize-weights

Run an optimization loop using Optuna to find optimal weight modification parameters. This command iteratively adjusts steering strength, layer positions, and other parameters to achieve a target metric value for any task type.

Basic Usage
python -m wisent optimize-weights --task TASK --model MODEL --output-dir DIR [OPTIONS]

Examples

Refusal Optimization
python -m wisent optimize-weights \
  --task refusal \
  --model meta-llama/Llama-3.1-8B-Instruct \
  --output-dir ./optimized_model/ \
  --target-metric compliance_rate \
  --target-value 0.95 \
  --trials 100
Benchmark Optimization
python -m wisent optimize-weights \
  --task arc_easy \
  --model meta-llama/Llama-3.1-8B-Instruct \
  --output-dir ./smart_model/ \
  --target-metric accuracy \
  --target-value 0.85 \
  --trials 200 \
  --early-stop
Personalization Optimization
python -m wisent optimize-weights \
  --task personalization \
  --trait "responds with technical precision and uses formal language" \
  --model meta-llama/Llama-3.1-8B-Instruct \
  --output-dir ./formal_model/ \
  --method titan \
  --num-directions 5 \
  --trials 150
Multi-Benchmark Optimization
python -m wisent optimize-weights \
  --task arc_easy,gsm8k,hellaswag \
  --model meta-llama/Llama-3.1-8B-Instruct \
  --output-dir ./multi_bench_model/ \
  --cap-pairs-per-benchmark 50 \
  --trials 300 \
  --checkpoint ./checkpoint.json \
  --s3-bucket wisent-optimization-results
With Custom Evaluator
python -m wisent optimize-weights \
  --task custom \
  --custom-evaluator ./my_eval.py:score_function \
  --model meta-llama/Llama-3.1-8B-Instruct \
  --output-dir ./custom_model/ \
  --target-metric score \
  --target-value 0.9

Arguments

Input Source (mutually exclusive)

ArgumentDescription
--taskTask: refusal, personalization, custom, benchmark name, or comma-separated benchmarks
--steering-vectorsPath to pre-computed steering vectors (skip generation)

Required Arguments

ArgumentDescription
--modelHuggingFace model name or path
--output-dirDirectory to save the optimized model

Optimization Target

ArgumentDefaultDescription
--target-metriccompliance_rateMetric to optimize (compliance_rate, refusal_rate, accuracy)
--target-value0.95Target value for the metric
--directionautoOptimization direction: auto, maximize, minimize

Optimization Parameters

ArgumentDefaultDescription
--trials300Number of optimization trials
--startup-trials10Random startup trials before TPE
--early-stopfalseStop early if target value is reached
--early-stop-patience10Stop if no improvement for N trials

Search Space

ArgumentDefaultDescription
--strength-range0.3,2.0Min,max range for projection strength
--max-weight-range0.5,3.0Min,max range for kernel max weight
--position-range0.3,0.7Min,max range for kernel peak position (0-1)
--num-pairs100Number of contrastive pairs to generate

Weight Modification Method

ArgumentDefaultDescription
--methoddirectionalMethod: directional, additive, titan, prism, pulse
--componentso_proj, down_projComponents to modify
--num-directions5Directions per layer for multi-direction methods

Checkpointing

ArgumentDefaultDescription
--checkpoint-Path to checkpoint file (saves/resumes optimization)
--checkpoint-interval5Save checkpoint every N trials
--s3-bucket-S3 bucket to upload results on completion

Related Commands

Stay in the loop. Never miss out.

Subscribe to our newsletter and unlock Wisent insights.