# Readout representation on language models
Examine the representation size by reconstructing texts from noised features. The variance of the noise is calibrated to realize specified correlation distance between the original and noised feature.

```
250610_readout_language/
├── README.md
├── sample_texts.py         # randomly sample texts from C4 dataset
├── save_features.py        # save features of the models
├── reconstruct.py          # reconstruct sentences from original and perturbed features
├── notebooks/              # notebooks to visualize results
└── configs/                # config files
```

## Config
All configuration parameters are defined in a single YAML file.

| Section         | Key                                      | Type              | Optional | Description |
|-----------------|-------------------------------------------|-------------------|----------|-------------|
| Model           | `model.pretrained`                        | string            | No       | HuggingFace model identifier |
| Dtype           | `dtype`                                   | string            | Yes      | Dtype for feature inversion. By default `float32`. `bf16` and `float16` are also available. |
| Experiment      | `exp_name`                                | string            | No       | Name of the experiment. Used in result directories. |
| Data            | `data.dataset_name`                       | string            | No       | Dataset name for organizing output. |
|                 | `data.text_dir`                          | string            | No       | Path to directory containing source texts. |
|                 | `data.text_names`                        | list of strings   | Yes      | List of text names (without extension). |
|                 | `data.text_names_path`                   | string            | Yes      | YAML file listing text names. Overrides `image_names` if provided. |
| Noise           | `noise.target_corr_dists`                 | list of floats    | No       | Target correlation distances between original and noised features. |
|                 | `noise.noise_seeds`                       | list of integers  | No       | Random seeds for noise generation. |
| Sample          | `max_length`                              | int               | No       | Max length of the samples. |
| Batch           | `batch_size`                              | int               | No       | Number of images optimized simultaneously on one GPU. |
| Pipeline        | `pipeline.num_iterations`                 | int               | No       | Number of optimization iterations. |
|                 | `pipeline.eval_interval`                  | int               | Yes      | Interval for evaluation. Default is `1`. |
|                 | `pipeline.wandb_log_interval`             | int               | Yes      | Logging interval to Weights & Biases. Default is `1`. | 
| Optimizer       | `optimizer.lr`                            | float             | No       | Learning rate. |
|                 | `optimizer.scheduler.start_factor`        | float             | Depends  | Scheduler-specific: start scaling factor. |
|                 | `optimizer.scheduler.end_factor`          | float             | Depends  | Scheduler-specific: end scaling factor. |
| Layers          | `layer_indices`                           | list of int       | No       | Layer indices to extract features. |
| Logging         | `wandb`                                   | bool              | Yes      | Enable Weights & Biases logging. |

---

### Notes

- **Per-layer execution**: All reconstructions are run independently for each layer specified in `layers`.
- **Distributed setup**: Supports multi-node and multi-GPU execution via shared storage, with JSON database tracking and file locking.


## Usage

### 0. Sample texts
```
python scripts/250610_readout_language/sample_texts.py --n 64 --output_dir data/texts/c4_sampled/
```
This script samples texts from HuggingFace datasets (by default, `allenai/c4`), and save them in a specified directory.

### 1. Run experiments
```
python scripts/250610_readout_language/reconstruct.py path/to/the/config.yaml
```
Args:
- config_path: path to the config
- device: device to perform the experiments, default `cuda`.

It will reconstrct sentences from original and perturbed features that have speficied correlation distances. As a result, it will perform reconstructions for roughly $N_{text} \times N_{layer} \times N_{distance} \times N_{noise\_seeds} \times N_{gen\_seeds}$, which can be a fairly large number.

#### Multi-node and multi-gpu support.
To perform experiments efficiently by leveraging computational resources, this scripts allow you to run the experiment with multiple nodes, as long as they share the storage. You can just run the script multiple times on each nodes. Also, to leverage multiple GPU, run the script multiple times with changing `--device` parameter. The simple trick is a json file that tracks which combination has done and which are pending. Please note it may not work under some file systems.

#### Output
Results will be saved in `output/readout_language/results/{model_alias}/{dataset_name}/{exp_name}/`

This directory has following structure: layer -> seed -> images

```
{exp_dir}/
├── experiment_db.json          # progress and summary of the results
├── experiment_db.lock          # lock file for the experiment_db.json
├── {layer}/
│   ├── corr_dist_0.2/
│   │   │   ├── noise_seed_0/
│   │   │   │   ├── {text_name}/
│   │   │   │   │   ├── history.csv
│   │   │   │   │   ├── summary.json
│   │   │   ├── noise_seed_1/
│   │   │   │   └── {text_name}/
│   │   │   │       └── ...
│   ├── corr_dist_0.4/
│   │   └── ...
│   ├── ...
│   └── corr_dist_0.0/
│           └── noise_seed_None/    # distance 0 = no noise = no noise seeds
│               └── {text_name}/
│                   └── ...
```


The `experiment_db.json` contains the status of experiments and the short summary of results. It stores the following contents for each experiment:
- `status`: pending, running, finished, or error
- `correlation_distance`: the correlation distance between the original feature and the target feature speficied in the config. Note that there might be an error between this value and the actual correlation distance.
- `true_target_*`: distance between the original and the target feature. If there's no postfix, it means the distance in the target layer.
- `target_recon_*`: distance between the original and reconstructed feature at the final iteration. If there's no postfix, it means the distance in the target layer.
- `true_recon_*`: distance between the target and reconstructed feature at the final iteration. If there's no postfix, it means the distance in the target layer.


If you only need quantitative results of the reconstructed images, looking up the `experiment_db.json` would suffice.
