# Code for ICLR 2026 Submission: Echoes as Anchors

This repository contains the source code for the ICLR 2026 submission, "Echoes as Anchors: Probabilistic Bounds and Attention Refocusing in LLM Reasoning".

## Overview

This project investigates the impact of the "Echo of Prompt" (EOP) behavior in Large Language Models (LLMs) on their reasoning capabilities. The core of the research is to determine whether the repetition of user questions within the reasoning process is a beneficial mechanism for reasoning or merely a template artifact from supervised fine-tuning (SFT).

The codebase includes implementations for:
1.  **Echo-Distilled SFT (ED-SFT) Data Preparation**: Scripts to process conversational data into a format suitable for fine-tuning models to adopt the "echo-then-reason" pattern.
2.  **Echoic Prompting (EP) Evaluation**: An inference-time, training-free strategy to re-ground models, implemented via a two-stage generation process.
3.  **MLP Probe Analysis**: A method to train a small MLP classifier as a probe to detect repetition patterns in the model's thinking process.
4.  **Probabilistic & Attention Analysis**: A suite of tools to compute the Echo Likelihood Gap ($\Delta\mathcal{L}$) and analyze attention patterns to understand the mechanisms behind EOP.

## Directory Structure

```
iclr2026_submission/
├── src/                      # Core source code
│   ├── evaluation/           # Evaluation scripts (harness, analysis)
│   ├── data_processing/      # Data preprocessing tools
│   └── utils/                # Utility functions
├── scripts/                  # Helper scripts to run experiments
├── configs/                  # Configuration files
├── train_mlp/                # MLP Probe experiment pipeline
├── eval_with_harness.py      # Main evaluation script
├── requirements.txt          # Python dependencies
└── README.md                 # This file
```

## Setup

### 1. Environment

We recommend using Conda to manage the environment.

```bash
# Create and activate conda environment
conda create -n iclr_submission python=3.9
conda activate iclr_submission

# Install dependencies
pip install -r requirements.txt
```

### 2. Configure Paths

Before running any experiments, you must configure the model and data paths. Edit the configuration file: `configs/training_config.yaml`.

Replace placeholder paths like `/path/to/your/models` with the actual absolute paths on your system. These paths are used by various analysis and evaluation scripts.

```yaml
# configs/training_config.yaml
# ...
  main_model_path: "/path/to/your/main_model/"
  embedding_model_path: "/path/to/your/embedding_model/"
  output_root: "/path/to/your/output_dir"
# ...
```

## Reproducing Paper Results

This project involves three main experimental workflows that correspond to the paper's contributions: ED-SFT, MLP Probe Analysis, and EP Evaluation.

### Workflow 1: Echo-Distilled SFT (ED-SFT) Data Generation

The paper demonstrates that fine-tuning with echo-infused data (ED-SFT) improves reasoning. We provide the script to prepare data for this fine-tuning process using a powerful teacher model.

**Core Idea: Data Distillation for Echo-then-Reason Pattern**

Instead of relying on existing conversational data, we use a data distillation approach. The script `src/data_processing/prepare_ed_sft_data.py` leverages a strong teacher model (e.g., GPT-4, or a powerful local model served via an OpenAI-compatible API) to generate high-quality Chain-of-Thought (CoT) reasoning traces that begin with a natural "Echo of Prompt" (EOP).

This is achieved through carefully designed prompts that guide the teacher model to first restate the question in a natural way before proceeding to the solution. The script offers several prompt variations (`natural`, `simple`, `minimal`) to generate a diverse dataset and avoid learning rigid templates.

**Step 1: Prepare ED-SFT Data**
Use the `prepare_ed_sft_data.py` script to generate the echo-infused CoT data from a base dataset like GSM8K.

```bash
# Point to your OpenAI-compatible API endpoint
export OPENAI_API_BASE="http://localhost:8000/v1"
export OPENAI_API_KEY="EMPTY" # Or your actual key

# Run the distillation script
python src/data_processing/prepare_ed_sft_data.py \
    --model "your_teacher_model_name" \
    --sample-size 7000 \
    --output /path/to/your/ed_sft_data.json
```
This script will load the GSM8K training set, generate a response for each question using the teacher model, and format the output into a `.json` file ready for fine-tuning.

**Step 2: Fine-tuning**
The generated `ed_sft_data.json` can be used with any standard SFT training library (we use Llama Factory's default full SFT parameter setting for deepseek-distill-llama-8b and qwen3 series models, but you can use any other library). The training process itself is standard and not included in this repository, as the primary contribution is the data preparation method and subsequent analysis.

### Workflow 2: MLP Probe Training and Analysis

This workflow trains a classifier to detect repetition, which is used in our analyses.

**Step 1: Preprocess Data**
Extract `(question, think_content)` pairs from your raw conversational data.

```bash
python train_mlp/preprocess_data.py /path/to/raw_data.jsonl /path/to/qt_pairs.jsonl
```

**Step 2: Label Data for Repetition**
Use a strong LLM (e.g., via OpenAI API) to label whether the `<think>` content repeats the question.

```bash
# Set your OpenAI API key if required by your endpoint
export OPENAI_API_KEY="your_api_key"

# Run labeling (requires a GPU for the embedding model)
CUDA_VISIBLE_DEVICES=0 python train_mlp/label_repeat.py /path/to/qt_pairs.jsonl /path/to/labeled_data.jsonl
```

**Step 3: Train the MLP Probe**
Train the MLP classifier on the labeled data.

```bash
CUDA_VISIBLE_DEVICES=0 python train_mlp/train_repeat_mlp.py /path/to/labeled_data.jsonl /path/to/trained_mlp.pt
```

### Workflow 3: Echoic Prompting (EP) and Other Evaluations

**Important Note**: For fair comparison with baseline methods, the main Echoic Prompting (EP) implementation uses the MI-PEAKS framework. The EP implementation can be found in the MI-PEAKS repository, specifically in:
- Main implementation: `MI-Peaks/src/applications/repeat.py`
- Evaluation scripts: `MI-Peaks/src/scripts/run_repeat_multi_budget.sh`

This implementation performs two-stage generation:
1. **Stage 1**: Generate initial response with a fixed token budget
2. **Stage 2**: Inject the question reminder and continue generation

The key parameters for EP are:
- `repeat_prompt`: "Let me reconsider the original question."
- `continuation_prompt`: "So now I know that"
- `token_budget`: Varies by dataset (e.g., 256-3072 for GSM8K)

**To reproduce Echoic Prompting (EP) results with MI-PEAKS:**
```bash
# Navigate to MI-PEAKS directory
cd /path/to/MI-Peaks/src/scripts

# Run EP evaluation with multiple token budgets
bash run_repeat_multi_budget.sh
```

**Alternative implementation (included in submission):**
We also provide a standalone two-stage evaluation script in this repository:

```bash
# Example for two_stage_echo mode
python src/evaluation/two_stage_eval.py \
    --model_path /path/to/your/base_model \
    --tasks gsm8k \
    --mode two_stage_echo \
    --output_dir /path/to/your/ep_results
```

**To evaluate standard models or fine-tuned checkpoints:**
```bash
# Set visible devices for evaluation
export CUDA_VISIBLE_DEVICES="0,1"

# Evaluate a model checkpoint
python eval_with_harness.py \
    --exp_dir /path/to/your/model_checkpoint_dir \
    --use_multi_gpu \
    --tensor_parallel_size 2
```

### Analysis Scripts

The `src/evaluation/` directory contains various scripts to reproduce the paper's analyses, such as:
-   `compare_trimmed_accuracy.py`: Computes the Echo Likelihood Gap ($\Delta\mathcal{L}$).
-   `attention_from_converted_refactored.py`: Computes attention metrics for the refocusing analysis.

Please refer to the docstrings within each script for detailed usage instructions.
