# AMML-2526 — Systematic Evaluation of Pre-trained Models and Training on a
Biased Low-Resource Dataset

This repository contains all code, data, and experiments for the *Advanced Methods in Machine Learning* coursework (2025–26).
It implements a two-part empirical study:

- **Task 1 — Representation analysis**
  Reverse-engineering three pretrained VAEs using latent probes, reconstructions, posterior diversity, and UMAP embeddings.

- **Task 2 — Bias & mitigation**
  Evaluating how class imbalance affects downstream digit classification when using a fixed VAE latent representation, and testing simple mitigation strategies.

The project follows a **reproducible research** layout: experiments are scripted, shared code is modularised, and all runs are deterministic via fixed random seeds.

---

## Project structure

```
amml-2526/
├── main.py
├── AMML_assignment_script.ipynb
├── pyproject.toml
├── uv.lock
├── .python-version
├── README.md
│
├── data/
│   ├── amml_model0_weights.pth
│   ├── amml_model1_weights.pth
│   └── amml_model2_weights.pth
│
├── experiments/
│   ├── task1/
│   │   ├── exp01_elbo.py
│   │   ├── exp02_recon.py
│   │   ├── exp03_probe.py
│   │   ├── exp04_umap.py
│   │   └── exp05_posterior_diversity.py
│   │
│   └── task2/
│       ├── exp01_bias_analysis.py
│       ├── exp02_class_weighting.py
│       └── exp03_oversampling.py
│
└── src/
    └── amml/
        ├── __init__.py
        ├── data.py
        ├── vae.py
        ├── probe.py
        ├── metrics.py
        ├── plots.py
        └── utils.py
```

---

## Directory overview

### `data/`
Contains the **three pretrained VAEs** supplied for the coursework:

- `amml_model0_weights.pth`
- `amml_model1_weights.pth`
- `amml_model2_weights.pth`

All models share the same architecture and training data but differ in training objectives or regularisation strategies.

---

### `src/amml/` — Core library

Reusable code shared across all experiments:

| File | Purpose |
|------|--------|
| `vae.py` | VAE architecture, encoder/decoder, latent sampling |
| `data.py` | Dataset loading, splits, class imbalance control |
| `probe.py` | Linear probes and downstream classifiers |
| `metrics.py` | ELBO, reconstruction error, balanced accuracy |
| `plots.py` | UMAPs, recon grids, recall bar plots |
| `utils.py` | Seeding, paths, and general helpers |

---

### `experiments/task1/` — Representation analysis

Each script probes a different aspect of the learned latent space:

| Script | What it measures |
|-------|-----------------|
| `exp01_elbo.py` | Evidence Lower Bound (ELBO) |
| `exp02_recon.py` | Reconstruction fidelity |
| `exp03_probe.py` | Linear probe performance |
| `exp04_umap.py` | Latent-space geometry (UMAP) |
| `exp05_posterior_diversity.py` | Posterior variance and spread |

Together, these experiments infer **what kind of representations** each pretrained VAE has learned.

---

### `experiments/task2/` — Bias and mitigation

All Task-2 experiments use a **fixed latent representation (Model 0)** to isolate the effect of dataset bias and mitigation strategies.

| Script | Purpose |
|-------|--------|
| `exp01_bias_analysis.py` | Baseline performance on imbalanced data |
| `exp02_class_weighting.py` | Loss re-weighting for minority classes |
| `exp03_oversampling.py` | Training set balancing via oversampling |

Each experiment is evaluated across **five random stratified splits** to estimate variability and robustness.


---

## Running the project

This project uses **uv** for environment and dependency management.

```bash
uv venv
source .venv/bin/activate
uv sync
```

Run Task 1 experiments:
```bash
uv run python experiments/task1/exp01_elbo.py
uv run python experiments/task1/exp02_recon.py
...
```

Run Task 2 experiments:
```bash
uv run python experiments/task2/exp01_bias_analysis.py
uv run python experiments/task2/exp02_class_weighting.py
uv run python experiments/task2/exp03_oversampling.py
```

All scripts save figures and metrics automatically to structured output directories.