Assignment 2: Ablation Archaeology

⏰ Due: Wednesday, March 4 | 15% of final grade

Systematically ablate the components of an E(3)-equivariant GNN to understand what makes various components of an architecture useful or not.

📄 Assignment description (PDF)

Overview

You'll work with a charged particle dynamics dataset and a working EGNN implementation. Your task: build an equivariance test, systematically remove or modify architectural components, propose an improvement, and analyze what breaks (and what doesn't).

Structure

├── assignment-02.pdf             # Assignment description
├── assignment-02-starter.ipynb   # Starter code with EGNN, data, and training loop
├── requirements.txt              # Python dependencies
└── report/
    ├── report.tex                # Your report (NeurIPS format)
    └── neurips_2025.sty          # LaTeX style file

Getting Started

1. Accept the assignment

Click the GitHub Classroom link shared by the instructor. This creates your own private copy of this repository.

2. Clone your repository

git clone https://github.com/bu-ds595/assignment02-YOUR_USERNAME.git
cd assignment02-YOUR_USERNAME

3. Install dependencies

pip install -r requirements.txt

4. Open the notebook

VS Code: Open the folder, then open assignment-02-starter.ipynb

JupyterLab: Run jupyter lab and open the notebook

Google Colab: Upload the notebook and run:

!pip install jax jaxlib flax optax

Deliverables

1. Code/Notebook — All ablation experiments, equivariance tests, training curves, and visualizations
2. Written Report — Maximum 3 pages in NeurIPS format (use report/report.tex)

Submitting

git add .
git commit -m "Complete assignment 2"
git push

You can push multiple times. Only the final version at the deadline will be graded.

Resources

• Assignment description (PDF)
• Lecture notes: Symmetry-Preserving Neural Networks
• Lecture slides: Symmetry-Preserving Neural Networks
• E(n) Equivariant Graph Neural Networks (Satorras et al., 2021) — the original EGNN paper
• A Recipe for Training Neural Networks (Karpathy, 2019) — practical advice for debugging and training