Phaedra: Learning High-Fidelity Discrete Tokenization for the Physical Science

TL;DR

This paper introduces Phaedra, a novel discrete tokenization method tailored for scientific images, enhancing the fidelity of physical and spectral properties in PDE-related data for improved deep learning applications.

Contribution

Phaedra is a new tokenization approach inspired by classical methods, specifically designed to better capture physical and spectral properties in scientific images, outperforming existing tokenizers.

Findings

Phaedra improves reconstruction accuracy across PDE datasets.

It demonstrates strong out-of-distribution generalization.

It effectively handles real-world Earth observation and weather data.

Abstract

Tokens are discrete representations that allow modern deep learning to scale by transforming high-dimensional data into sequences that can be efficiently learned, generated, and generalized to new tasks. These have become foundational for image and video generation and, more recently, physical simulation. As existing tokenizers are designed for the explicit requirements of realistic visual perception of images, it is necessary to ask whether these approaches are optimal for scientific images, which exhibit a large dynamic range and require token embeddings to retain physical and spectral properties. In this work, we investigate the accuracy of a suite of image tokenizers across a range of metrics designed to measure the fidelity of PDE properties in both physical and spectral space. Based on the observation that these struggle to capture both fine details and precise magnitudes, we propose Phaedra, inspired by classical shape-gain quantization and proper orthogonal decomposition. We demonstrate that Phaedra consistently improves reconstruction across a range of PDE datasets. Additionally, our results show strong out-of-distribution generalization capabilities to three tasks of increasing complexity, namely known PDEs with different conditions, unknown PDEs, and real-world Earth observation and weather data.