CATS

This repository contains the core implementation for CATS: Category-Aware Token-level Steering for Training-Free Redundancy Reduction in Large Reasoning Models . The project focuses on three main components:

1. Hidden State Extraction - Extract hidden states from transformer models
2. Layer Intervention - Perform targeted interventions on specific neural layers
3. Redundancy Classification - Classify model outputs for redundancy patterns

Features

• Extract hidden states from transformer models at inference time
• Intervene on specific layers during model generation
• Classify model responses for different types of redundancy
• Parallel processing support for batch operations
• Configurable model paths and data sources

Installation

1. Clone this repository:

git clone <repository-url>
cd CATS

2. Install dependencies:

pip install -r requirements.txt

3. Configure settings:

cp config/config.template.json config/config.json
# Edit config/config.json with your specific paths and settings

Usage

Hidden State Extraction

Extract hidden states from model responses:

# Simple mode - single text input
python src/hidden_states/extract_hidden_states.py \
    --mode simple \
    --text "What is the future of artificial intelligence?"

# Batch mode - process file
python src/hidden_states/extract_hidden_states.py \
    --mode batch \
    --input_file data/input/questions.json \
    --output_dir data/output/hidden_states \
    --batch_size 8

Layer Intervention

Perform interventions on specific layers:

# Single layer intervention
python src/intervention/layer_intervention.py \
    --mode single \
    --layer_idx 12 \
    --input_text "What is 2 + 2?" \
    --intervention_file data/intervention_vectors.pkl

# All layers intervention
python src/intervention/layer_intervention.py \
    --mode all \
    --input_text "What is 2 + 2?" \
    --intervention_file data/intervention_vectors.pkl \
    --output_file results/intervention_results.json

Redundancy Classification

Classify responses for redundancy patterns:

python src/classification/redundancy_classifier.py \
    --original_file data/input/original_data.json \
    --llm_responses_dir data/output/llm_responses \
    --output_file results/classification_results.json \
    --max_workers 4

Configuration

Edit config/config.json to set:

• Model paths: Specify your transformer model location
• Data directories: Set input/output data paths
• API configurations: Configure API keys (use placeholder values)
• Processing parameters: Batch sizes, worker counts, etc.

Example Configuration

{
  "model": {
    "model_path": "/path/to/your/model",
    "device_map": "auto",
    "torch_dtype": "bfloat16"
  },
  "data": {
    "input_data_dir": "./data/input",
    "output_data_dir": "./data/output"
  },
  "processing": {
    "batch_size": 8,
    "max_workers": 4
  },
  "api": {
    "api_keys": ["your-api-key-here"]
  }
}

Project Structure

├── src/
│   ├── hidden_states/     # Hidden state extraction modules
│   │   └── extract_hidden_states.py
│   ├── intervention/      # Layer intervention methods
│   │   └── layer_intervention.py
│   ├── classification/    # Redundancy classification
│   │   └── redundancy_classifier.py
│   └── utils/            # Utility functions and prompts
│       ├── prompts.py
│       └── data_utils.py
├── config/               # Configuration files
│   └── config.template.json
├── tests/               # Test scripts
├── data/               # Data directory (create as needed)
│   ├── input/
│   └── output/
├── requirements.txt     # Python dependencies
└── README.md           # This file

Data Format

Input Data Format

Questions file (data/input/questions.json):

[
  {
    "idx": 0,
    "question": "What is 2 + 2?",
    "answer": "4"
  }
]

Output Data Format

Hidden states are saved as pickle files containing layer-wise vectors. Classification results are saved as JSON with redundancy type annotations.

Redundancy Types

The system classifies responses into the following redundancy categories:

Destructive Redundancy:

1. Logical Drift
2. Hallucination Amplification
3. Internal Contradiction

Non-Destructive Redundancy: 4. Repetition 5. Over-Elaboratio 6. Over-Cautiousness

License

This project is licensed under the MIT License - see the LICENSE file for details.

Notes

• Ensure you have sufficient GPU memory for model loading
• API keys should be kept secure and not committed to version control
• For large datasets, consider using the batch processing mode
• Monitor memory usage during hidden state extraction