Inconsistent Pre-trained Baseline Performance

[ArshKA]

Hello, I was testing the baseline model checkpoints and I'm not getting close to the reported numbers for some models. For example, viscoelastic_instability, FNO and TFNO have similar metrics compared to the reported value, but both UNets are very different. Is there anything wrong in the way I'm evaluating these models? Attached a minimal code snippet below

model: polymathic-ai/FNO-viscoelastic_instability
Mean Normalized VRMSE: 0.712697 (Std Dev: 0.033637)
Mean Denormalized VRMSE: 0.706040 (Std Dev: 0.033151)
Reported Number: 0.7212

model: polymathic-ai/TFNO-viscoelastic_instability
Mean Normalized VRMSE: 0.708173 (Std Dev: 0.030011)
Mean Denormalized VRMSE: 0.702301 (Std Dev: 0.028892)
Reported Number:  0.7102

model: polymathic-ai/UNetConvNext-viscoelastic_instability
Mean Normalized VRMSE: 0.946716 (Std Dev: 0.034722)
Mean Denormalized VRMSE: 0.937711 (Std Dev: 0.029086)
Reported Number:  0.2499

model: polymathic-ai/UNetClassic-viscoelastic_instability
Mean Normalized VRMSE: 3.949996 (Std Dev: 1.488076)
Mean Denormalized VRMSE: 3.870942 (Std Dev: 1.455580)
Reported Number: 0.4185

Here's the code I'm using:

import numpy as np
import torch
from einops import rearrange
from tqdm import tqdm
import random

from the_well.benchmark.metrics import VRMSE
from the_well.data import WellDataset
from the_well.utils.download import well_download
from the_well.benchmark.models import FNO, UNetConvNext, UNetClassic, TFNO
from the_well.data.normalization import ZScoreNormalization

base_path = '/data/the_well/datasets'
device = torch.device("cuda")
dataset = WellDataset(
    well_base_path=base_path,
    well_dataset_name="viscoelastic_instability",
    well_split_name="test",
    n_steps_input=4,
    n_steps_output=1,
    use_normalization=True,
    normalization_type=ZScoreNormalization,
)

model_path = 'polymathic-ai/UNetClassic-viscoelastic_instability'
model = UNetClassic.from_pretrained(model_path).to(device)
model.eval()

n_samples_per_batch = 8
n_batches = 100
total_samples = n_samples_per_batch * n_batches

all_batch_mean_normalized_errors = []
all_batch_mean_denormalized_errors = []


for batch_idx in tqdm(range(n_batches), desc="Evaluating Batches"):
    input_batches = []
    output_batches = []
    for _ in range(n_samples_per_batch):
        random_idx = random.randrange(len(dataset))
        sample = dataset[random_idx]
        input_batches.append(sample['input_fields'].to(device))
        output_batches.append(sample['output_fields'].to(device))

    input_batch = torch.stack(input_batches)  # Shape: (B, Ti, Lx, Ly, F)
    output_batch = torch.stack(output_batches) # Shape: (B, To, Lx, Ly, F)

    input_batch = rearrange(input_batch, "B Ti Lx Ly F -> B (Ti F) Lx Ly")

    with torch.no_grad():
        pred_batch = model(input_batch) # Shape: (B, (To*F), Lx, Ly)

    pred_batch = rearrange(pred_batch, "B (Tp F) Lx Ly -> B Tp Lx Ly F", Tp=dataset.n_steps_output) # Shape: (B, Tp, Lx, Ly, F)

    normalized_errors_per_sample = VRMSE.eval(pred_batch, output_batch, dataset.metadata)
    mean_normalized_err = normalized_errors_per_sample.mean().item()
    all_batch_mean_normalized_errors.append(mean_normalized_err)

    denormalized_pred_batch = dataset.norm.denormalize_flattened(pred_batch, mode="variable")
    denormalized_output_batch = dataset.norm.denormalize_flattened(output_batch, mode="variable")
    denormalized_errors_per_sample = VRMSE.eval(denormalized_pred_batch, denormalized_output_batch, dataset.metadata)
    mean_denormalized_err = denormalized_errors_per_sample.mean().item()
    all_batch_mean_denormalized_errors.append(mean_denormalized_err)


final_mean_normalized = np.mean(all_batch_mean_normalized_errors)
final_std_normalized = np.std(all_batch_mean_normalized_errors)
final_mean_denormalized = np.mean(all_batch_mean_denormalized_errors)
final_std_denormalized = np.std(all_batch_mean_denormalized_errors)

print("\n--- Overall Evaluation Results ---")
print(f"Total samples evaluated: {total_samples} ({n_batches} batches of {n_samples_per_batch} samples)")
print(f"Mean Normalized VRMSE: {final_mean_normalized:.6f} (Std Dev: {final_std_normalized:.6f})")
print(f"Mean Denormalized VRMSE: {final_mean_denormalized:.6f} (Std Dev: {final_std_denormalized:.6f})")

[mikemccabe210]

Thanks ArshKA and Ryno. This is related to a bug fix in the test loader where validation was not being applied in the same way as train and validation. We're in the process of updating the paper and metrics to reflect the fix and should have everything fixed up by early next week. We'll post another notification once that's done.

Full release notes for the latest version can be found here though the normalization and the outlier trajectory in acoustic_scattering_inclusions should be the only changes impacting metrics.

[jacklishufan]

HI. I'm wondering if the authors can provide an update on the process. Specifically, I'm interested to know which of the following is the case:

1. The released training and evaluation code is broken and as such the results cannot be reproduced due to a fixable bug.
2. Some baseline numbers reported in the original paper are incorrect and needs to be updated.

We plan to use this code base for an ongoing project. I would appreciate it if authors can provide some insights.

[LTMeyer]

Hi @jacklishufan,

The main issue was that normalization was not applied consistently during training and validation.

1. The released training and evaluation code is broken and as such the results cannot be reproduced due to a fixable bug.

The master branch has already resolved this issue. You can use the master branch for training and evaluation in your own project.

2. Some baseline numbers reported in the original paper are incorrect and needs to be updated.

Due to the issue mentioned above there are indeed some discrepancies with the numbers reported in the original paper. This however only concerns a few datasets. It is in the process to be fixed. It only requires to re-run the evaluation.

[hassony2]

Hi !

I was wondering if the evaluation has been re-run by someone following the fix at head. It would be great to have the final numbers for future comparisons :)

Best,

[till-m]

Hi everyone,

I noticed that these scores are still not updated, which is concerning as there are already pre-prints using them for comparison, e.g. this one, with only a small mention of the issue in Appendix D. It would be good to fix the problem now before these errors propagate further.

If this is a matter of compute, I might be able to run tests for some problems on the cluster of my university.

[Ryno666]

谢谢，您的邮件我已收到！

[mikemccabe210]

Mostly a human-power issue, though that's been sorted out and we'll have an update out in the next few weeks.

[till-m]

@mikemccabe210 That's excellent, thanks!

I'm wondering if test-time normalization is only part of the story here. Could there be other differences in how the baseline model weights on HF were trained?
I'm asking because, when I trained a ConvNextUnet from scratch myself (using the-well v1.1.0), it significantly outperforms the pre-trained weights, and even scores slightly better than what's reported in the paper. This makes me think that there might be additional discrepancies beyond just normalization settings.

--------------------------------------------------------------------------------
PRE-TRAINED UNetConvNext (polymathic-ai)
--------------------------------------------------------------------------------
Mean MSE:   9.373346e-01 ± 2.957979e-01
Mean VRMSE: 0.946770 ± 0.034963

--------------------------------------------------------------------------------
SELF-TRAINED UNetConvNext (epoch 78)
--------------------------------------------------------------------------------
Mean MSE:   9.593404e-02 ± 1.134732e-01
Mean VRMSE: 0.225992 ± 0.047241

Comparison Code

To run, use the model weights from here.

"""
Compare pre-trained vs self-trained UNetConvNext models on viscoelastic instability.
Reports raw MSE and VRMSE loss values.
"""

import numpy as np
import torch
from tqdm import tqdm
import random

from the_well.data import WellDataset
from the_well.benchmark.models import UNetConvNext
from the_well.data.normalization import ZScoreNormalization
from einops import rearrange

# Configuration
BASE_PATH = 'data/the-well/datasets'
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
N_SAMPLES_PER_BATCH = 8
N_BATCHES = 100
TOTAL_SAMPLES = N_SAMPLES_PER_BATCH * N_BATCHES

print(f"Using device: {DEVICE}")
print(f"Total samples to evaluate: {TOTAL_SAMPLES} ({N_BATCHES} batches of {N_SAMPLES_PER_BATCH})")
print("="*80)

# Load dataset
print("\nLoading viscoelastic_instability test dataset...")
dataset = WellDataset(
    well_base_path=BASE_PATH,
    well_dataset_name="viscoelastic_instability",
    well_split_name="test",
    n_steps_input=4,
    n_steps_output=1,
    use_normalization=True,
    normalization_type=ZScoreNormalization,
)
print(f"Dataset loaded: {len(dataset)} samples")
print(f"Spatial resolution: {dataset.metadata.spatial_resolution}")
print(f"Number of fields: {dataset.metadata.n_fields}")

# Load pre-trained model
print("\nLoading pre-trained UNetConvNext...")
model_pretrained = UNetConvNext.from_pretrained(
    "polymathic-ai/UNetConvNext-viscoelastic_instability"
).to(DEVICE)
model_pretrained.eval()
print("Pre-trained model loaded")

# Load self-trained model
print("\nLoading self-trained UNetConvNext...")
# using .from_pretrained instead of a manual __init__ call 
model_selftrained = UNetConvNext.from_pretrained(
    "polymathic-ai/UNetConvNext-viscoelastic_instability"
)

# checkpoint available from https://drive.switch.ch/index.php/s/I8DFR2x0QBS7KdY
checkpoint = torch.load("self_trained_cnunet.pt", map_location='cpu')
model_selftrained.load_state_dict(checkpoint['model_state_dict'])
model_selftrained = model_selftrained.to(DEVICE)
model_selftrained.eval()
print(f"Self-trained model loaded from epoch {checkpoint['epoch']}")

# Storage for batch-wise metrics
pretrained_mse_batches = []
pretrained_vrmse_batches = []
selftrained_mse_batches = []
selftrained_vrmse_batches = []

# Evaluation loop
print("\n" + "="*80)
print("Starting evaluation...")
print("="*80)

for batch_idx in tqdm(range(N_BATCHES), desc="Evaluating batches"):
    # Collect samples for this batch
    input_batches = []
    output_batches = []

    for _ in range(N_SAMPLES_PER_BATCH):
        random_idx = random.randrange(len(dataset))
        sample = dataset[random_idx]
        input_batches.append(sample['input_fields'].to(DEVICE))
        output_batches.append(sample['output_fields'].to(DEVICE))

    # Stack into batches
    input_batch = torch.stack(input_batches)   # (B, Ti, Lx, Ly, F)
    output_batch = torch.stack(output_batches)  # (B, To, Lx, Ly, F)

    # Rearrange for model input: (B, Ti, Lx, Ly, F) -> (B, Ti*F, Lx, Ly)
    input_batch_formatted = rearrange(input_batch, "B Ti Lx Ly F -> B (Ti F) Lx Ly")

    with torch.no_grad():
        # Get predictions from both models
        pred_pretrained = model_pretrained(input_batch_formatted)  # (B, To*F, Lx, Ly)
        pred_selftrained = model_selftrained(input_batch_formatted)

        # Rearrange predictions: (B, To*F, Lx, Ly) -> (B, To, Lx, Ly, F)
        pred_pretrained = rearrange(
            pred_pretrained,
            "B (To F) Lx Ly -> B To Lx Ly F",
            To=dataset.n_steps_output
        )
        pred_selftrained = rearrange(
            pred_selftrained,
            "B (To F) Lx Ly -> B To Lx Ly F",
            To=dataset.n_steps_output
        )

        # Compute MSE (mean over spatial dimensions, averaged over batch/time/fields)
        pretrained_mse = torch.mean((pred_pretrained - output_batch) ** 2, dim=(2, 3))  # (B, To, F)
        selftrained_mse = torch.mean((pred_selftrained - output_batch) ** 2, dim=(2, 3))

        # Average MSE over batch/time/fields for this batch
        pretrained_mse_batches.append(pretrained_mse.mean().item())
        selftrained_mse_batches.append(selftrained_mse.mean().item())

        # Compute VRMSE: sqrt(MSE / variance)
        gt_variance = torch.var(output_batch, dim=(2, 3)) + 1e-7  # (B, To, F)
        pretrained_vrmse = torch.sqrt(pretrained_mse / gt_variance)
        selftrained_vrmse = torch.sqrt(selftrained_mse / gt_variance)

        # Average VRMSE over batch/time/fields for this batch
        pretrained_vrmse_batches.append(pretrained_vrmse.mean().item())
        selftrained_vrmse_batches.append(selftrained_vrmse.mean().item())

# Compute final statistics
pretrained_mse_mean = np.mean(pretrained_mse_batches)
pretrained_mse_std = np.std(pretrained_mse_batches)
pretrained_vrmse_mean = np.mean(pretrained_vrmse_batches)
pretrained_vrmse_std = np.std(pretrained_vrmse_batches)

selftrained_mse_mean = np.mean(selftrained_mse_batches)
selftrained_mse_std = np.std(selftrained_mse_batches)
selftrained_vrmse_mean = np.mean(selftrained_vrmse_batches)
selftrained_vrmse_std = np.std(selftrained_vrmse_batches)

# Print results
print("\n" + "="*80)
print("FINAL RESULTS - MODEL COMPARISON")
print("="*80)
print(f"\nTotal samples evaluated: {TOTAL_SAMPLES}")
print("\n" + "-"*80)
print("PRE-TRAINED UNetConvNext (polymathic-ai)")
print("-"*80)
print(f"Mean MSE:   {pretrained_mse_mean:.6e} ± {pretrained_mse_std:.6e}")
print(f"Mean VRMSE: {pretrained_vrmse_mean:.6f} ± {pretrained_vrmse_std:.6f}")

print("\n" + "-"*80)
print("SELF-TRAINED UNetConvNext (epoch {})".format(checkpoint['epoch']))
print("-"*80)
print(f"Mean MSE:   {selftrained_mse_mean:.6e} ± {selftrained_mse_std:.6e}")
print(f"Mean VRMSE: {selftrained_vrmse_mean:.6f} ± {selftrained_vrmse_std:.6f}")

print("\n" + "="*80)
print("COMPARISON")
print("="*80)

# MSE comparison
if pretrained_mse_mean < selftrained_mse_mean:
    mse_improvement = ((selftrained_mse_mean - pretrained_mse_mean) / selftrained_mse_mean) * 100
    mse_better = "PRE-TRAINED"
else:
    mse_improvement = ((pretrained_mse_mean - selftrained_mse_mean) / pretrained_mse_mean) * 100
    mse_better = "SELF-TRAINED"

print(f"\nMSE:   {mse_better} is better by {mse_improvement:.2f}%")

# VRMSE comparison
if pretrained_vrmse_mean < selftrained_vrmse_mean:
    vrmse_improvement = ((selftrained_vrmse_mean - pretrained_vrmse_mean) / selftrained_vrmse_mean) * 100
    vrmse_better = "PRE-TRAINED"
else:
    vrmse_improvement = ((pretrained_vrmse_mean - selftrained_vrmse_mean) / pretrained_vrmse_mean) * 100
    vrmse_better = "SELF-TRAINED"

print(f"VRMSE: {vrmse_better} is better by {vrmse_improvement:.2f}%")
print("="*80)

Full Output

$ uv run compare_unet_models.py 
.venv/lib/python3.13/site-packages/torch/cuda/__init__.py:63: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
  import pynvml  # type: ignore[import]
.venv/lib/python3.13/site-packages/pydantic/_internal/_generate_schema.py:2249: UnsupportedFieldAttributeWarning: The 'repr' attribute with value False was provided to the `Field()` function, which has no effect in the context it was used. 'repr' is field-specific metadata, and can only be attached to a model field using `Annotated` metadata or by assignment. This may have happened because an `Annotated` type alias using the `type` statement was used, or if the `Field()` function was attached to a single member of a union type.
  warnings.warn(
.venv/lib/python3.13/site-packages/pydantic/_internal/_generate_schema.py:2249: UnsupportedFieldAttributeWarning: The 'frozen' attribute with value True was provided to the `Field()` function, which has no effect in the context it was used. 'frozen' is field-specific metadata, and can only be attached to a model field using `Annotated` metadata or by assignment. This may have happened because an `Annotated` type alias using the `type` statement was used, or if the `Field()` function was attached to a single member of a union type.
  warnings.warn(
.venv/lib/python3.13/site-packages/timm/models/layers/__init__.py:48: FutureWarning: Importing from timm.models.layers is deprecated, please import via timm.layers
  warnings.warn(f"Importing from {__name__} is deprecated, please import via timm.layers", FutureWarning)
Using device: cuda
Total samples to evaluate: 800 (100 batches of 8)
================================================================================

Loading viscoelastic_instability test dataset...
Dataset loaded: 592 samples
Spatial resolution: (512, 512)
Number of fields: 8

Loading pre-trained UNetConvNext...
Pre-trained model loaded

Loading self-trained UNetConvNext...
Self-trained model loaded from epoch 78

================================================================================
Starting evaluation...
================================================================================
Evaluating batches: 100%|███████████████████████████████████████████████████████████████████████████████████████| 100/100 [01:49<00:00,  1.09s/it]

================================================================================
FINAL RESULTS - MODEL COMPARISON
================================================================================

Total samples evaluated: 800

--------------------------------------------------------------------------------
PRE-TRAINED UNetConvNext (polymathic-ai)
--------------------------------------------------------------------------------
Mean MSE:   9.373346e-01 ± 2.957979e-01
Mean VRMSE: 0.946770 ± 0.034963

--------------------------------------------------------------------------------
SELF-TRAINED UNetConvNext (epoch 78)
--------------------------------------------------------------------------------
Mean MSE:   9.593404e-02 ± 1.134732e-01
Mean VRMSE: 0.225992 ± 0.047241

================================================================================
COMPARISON
================================================================================

MSE:   SELF-TRAINED is better by 89.77%
VRMSE: SELF-TRAINED is better by 76.13%
================================================================================

Visual Comparison

[mikemccabe210]

Thanks @till-m, that's very helpful. And yeah, there are two main issues we identified. One is the normalization flag changing, the second is that some models were trained before we normalized the orientation of some datasets (ie, a pre-release version once had dataset A is stored in [x, y, z], dataset B is [z, y, x]), so the huggingface models are reflecting the old orientation. In general, visco is a bit of an outlier data-wise since that dataset wasn't amenable to ensuring train/valid/test are all drawn from the same distribution, so smaller changes there might have a larger impact than most.

To give a bit of a sense of what's going on at the moment - right now, one of our new team members is going through the old checkpoints and seeing which reproduce the exact numbers (within a slight margin of error for numerics changes over software versions) and in the cases where it doesn't, retraining the model to confirm it still ends up in the generally vicinity. The timeline looks kind of like:

• We'll be releasing a regular software update this week with some new features. This is unrelated to this thread, though originally we wanted to to sync the two matters, but realized that's not going to happen.
• Once the analysis of the existing checkpoints is complete, we'll post those results here as "unofficial" numbers for reference. This should be this week or next week.
• Afterwards, we'll queue up re-runs of everything at the "winning" hyperparameter configuration on the latest version of the code with both our team's changes and publicly contributed changes. These retrained checkpoints will be released to huggingface and we'll update the paper text to reflect these as the "official" numbers to make sure public checkpoints, numbers, and code match exactly up to hardware differences/stochasticity. This could still take another month or so. Realistically - these will probably all run while everyone using the cluster at Flatiron is at NeurIPS.