# The Platonic Transformer: A Solid Choice for Equivariance

This codebase contains the source for the Platonic Transformer, submitted to ICLR 2026.

![Platonic Transformer Architecture](assets/platoformer_architecture.svg)

## Organization

```
|--datasets/
|--models/
|----baseline/
|------esen/
|------faformer/
|----platoformer/
|----rapidash/
|--requirements.txt
|--README.md
|--utils.py
|--main_*.py
```

The main directory provides scripts to run the models for different tasks on the datasets used in the paper. The Platonic Transformer architecture is implemented in `models/platoformer/` alongside baselines taken from their respective open-source sources. We do not include Geometric Hyena (Moskalev et al., 2025) as its code was obtained through private correspondence with the authors.

## Setup

We recommend using [`uv`](https://docs.astral.sh/uv/) to set up the environment:

```bash
$ uv venv -p 3.12.4
$ source .venv/bin/activate
$ uv pip install -r requirements.txt
```

To active the environment,

```bash
$ source .venv/bin/activate
```

> We recommend using CUDA 12.1 and above on your GPUs with at least 40GiB of VRAM (eg: A100, RTX 4090, H200, etc).

## Training

We provide separate training scripts for each task, separated into **Non-equivariant** and **Equivariant**.

| File                   | Task            | Dataset      | Task                          |
|------------------------|-----------------|--------------|-------------------------------|
| `main_cifar10.py`      | Non-equivariant | CIFAR-10     | Image classification          |
| `main_scanobjectnn.py` | Non-equivariant | ScanObjectNN | 3D point cloud classification |
| `main_omol.py`         | Equivariant     | OMol25       | Molecular property prediction |
| `main_qm9_regr.py`     | Equivariant     | QM9          | Molecular property prediction |

Within each file, we provide comments on different arguments to pass via CLI. To reproduce results from our paper, refer to the Appendix in the submitted PDF for details on exact hyperparameters for each task/dataset.

#### Additional Details

- OMol25 must be installed beforehand locally to train on it. We use the Neutral split with 34M molecules for training, 27K molecules for validation, and 2M molecules for testing. Installation instructions are provided [on HuggingFace](https://huggingface.co/facebook/OMol25). To set the dataset path in `main_omol.py`, you can use the `--data_dir` flag.

- ScanObjectNN must be installed beforehand locally to train on it. Download instructions are provided on [the website](https://hkust-vgd.github.io/scanobjectnn/).

- We use *Weights and Biases* (WandB) to monitor experiments. You must have logged in via CLI. Instructions to set up WandB via CLI can be found [here](https://docs.wandb.ai/ref/cli/wandb-login/).