Efficiently Exporting AlphaEarth Embeddings from GEE

Google recently released the AlphaEarth Foundations annual Satellite Embeddings dataset, a 64-dimensional vector representation of Earth’s surface that incorporates spectral reflectance, radar backscatter, climate, and topography in an analysis-ready format. Currently, the dataset is only available through Earth Engine, so if you want to work with it outside the platform, you’ll need to export it first¹.

Earth Engine stores the embeddings in 64-bit floating point images. Exporting as-is might be manageable for small areas, but 64 bands at 10m resolution and 64 bits per pixel adds up quickly – the full 7-year time series of annual embeddings across CONUS at 10m resolution would be pushing a petabyte, uncompressed.

You could save some storage and RAM by dropping precision to 32 bits or scaling and rounding to 16-bit integers, but there’s a more efficient, lossless way to store the embeddings in just 8 bits via quantization.

Byte embeddings

To understand why it’s possible to store the embeddings with 8x fewer bits without losing any information, we need to dive a little deeper and take a look at the Brown et. al. preprint that accompanies the dataset. Specifically, this sentence:

… we quantize the 32-bit floating point embeddings generated by AEF to 8 bits, resulting in an 4x reduction in storage with negligible impact on performance …

Quantization is commonly used in deep learning to generate smaller models with lower computational overhead by reducing the precision of their weights via truncating or scaling and rounding to a smaller data type². AlphaEarth’s embeddings are effectively just learned model weights in a 64D latent space, so it makes sense that the authors chose to quantize, from the original 32-bit outputs to 8 bits for the final product.

But if the embeddings were quantized to 8 bits, why are they stored in 64-bit floats in Earth Engine? De-quantization.

Quantization sacrifices a small amount of precision in exchange for efficient storage, but it also sacrifices some usability. The original model outputs are floating-point, normalized, unit-length vectors that can be fed directly into most ML algorithms. The AlphaEarth quantization process removes those properties by square-root transforming, scaling, and rounding the embeddings into signed 8-bit integers. To preserve the embeddings as an analysis-ready dataset, the Earth Engine team clearly decided to de-quantize the 8-bit embeddings into 64 bits³, at the cost of storage and memory.

It’s worth pointing out that unlike the raw un-quantized model outputs, the de-quantized embeddings in Earth Engine still only contain 8 bits of data, which becomes obvious when you look at the number of repeated values within the embeddings:

So, the embeddings stored in Earth Engine are 8 bits of data, de-quantized to 64 bits. If we want to export them to the smallest possible size, we just need to re-quantize them back to 8 bits.

Re-quantizing

The supplementary materials for Brown et. al. helpfully provide some pseudo-code and parameters for their quantization process, which is easy to adapt to Earth Engine.

The quantization is essentially just a square root transformation followed by scaling and rounding to a signed int8:

POWER = 2.0
SCALE = 127.5
MIN_VALUE = -127.0
MAX_VALUE = 127.0

def quantize(img: ee.Image) -> ee.Image:
    """Quantize 64-bit AlphaEarth embeddings to 8-bit integers."""
    sat = img.abs().pow(1 / POWER).multiply(img.signum())
    return sat.multiply(SCALE).clamp(MIN_VALUE, MAX_VALUE).int8()

For testing, we’ll also want to be able to de-quantize in Earth Engine, which just reverses the process:

def dequantize(img: ee.Image, ) -> ee.Image:
    """De-quantize 8-bit AlphaEarth embeddings to 64-bit."""
    rescaled = img.double().divide(SCALE)
    return rescaled.abs().pow(POWER).multiply(rescaled.signum())

To confirm that this is a lossless transformation, we can use those functions to quantize an image to 8-bit and de-quantize back to 64-bit. Any differences between the original de-quantized image and our round-trip de-quantized image would indicate loss in precision, so I ran a reducer across the image to look for errors:

img = ee.Image("GOOGLE/SATELLITE_EMBEDDING/V1/ANNUAL/x2codd2zmw954ofn6")

# Calculate error in round-trip quantization
quantized = quantize(img)
dequantized = dequantize(quantized)
error = img.subtract(dequantized).abs()

# Get the largest error across all bands of the image
max_error = error.reduceRegion(
    reducer=ee.Reducer.max(),
    geometry=img.geometry(),
    scale=10,
    maxPixels=1e13,
).values().reduce(ee.Reducer.max()).getInfo()

print(max_error)

0

Zero error confirms that this is a lossless conversion.

To quantize or not to quantize

That’s everything needed to export the quantized embeddings from Earth Engine into 8-bit GeoTIFFs, but is it actually worth the effort?

Quantization theoretically reduces storage size by 8x. How much impact does it have in practice, with a compression method like LZW? In my experience, the actual storage savings over the original de-quantized images are closer to ~2x.

$ du -h ./embeddings_*.tif

657M	./embeddings_float64.tif
363M	./embeddings_int8.tif

Whether the data is de-quantized or not, the original quantization already reduced the precision of the embeddings, and the relatively small number of unique values leads to high compression ratios, regardless of data type.

Once you bring the quantized data into memory, you might actually see an 8x reduction in RAM, if your analysis works with quantized data. For my purposes of extracting embeddings to plot locations and slicing small windows out of larger embeddings, I can store the bulk of the data quantized and then de-quantize on-the-fly at manageable scales. But if you’re going to de-quantize everything every time you run analysis, you’re probably better off skipping the quantization.

Like most optimizations, quantizing AlphaEarth embeddings to 8 bits is a tradeoff – in this case, you’re trading a smaller storage and memory footprint for more processing if and when you need to de-quantize. Whether that’s a worthwhile trade will depend on your hardware and your use case.

Here’s the function I use to de-quantize locally after loading quantized GeoTIFFs with Xarray. Note that I’m de-quantizing to 32 bits locally, since it saves some space and you won’t gain anything from the extra “precision” of 64 bits. I also encoded NoData as -128, since that’s unused in the quantized range.

def dequantize(da: xr.DataArray) -> xr.DataArray:
    """De-quantize 8-bit AlphaEarth embeddings to 32-bit."""
    masked = da.astype("float32").where(da != -128)
    rescaled = masked / SCALE
    return abs(rescaled) ** POWER * np.sign(rescaled)