Question about reproducing bert OSM tag pretraining

[xxn967]

Hi, thank you for releasing the code, models, and data.

I am trying to reproduce the BERT OSM tag pretraining experiment using:

• nlp_pretraining/osm_tag_embed_from_scratch.py
• the released dataset relevant_tags_pairs_dataset_parquet_sharded_20_rep.tar.gz
• the released initial model bert-144-l6-reset-custom-tokenizer

However, I am still seeing a significant gap compared with the released BERT checkpoint, so I would like to confirm some information and details.

From the released training artifacts and script, I understand the main setup to be roughly:

• CachedMultipleNegativesSymmetricRankingLoss
• num_train_epochs = 4
• learning_rate = 1e-3
• warmup_ratio = 0.1
• fp16 = True
• BatchSamplers.NO_DUPLICATES

The main difference I can clearly identify on my side is the batch-related setting, but I am not sure whether this alone explains the discrepancy.
For batchsize and minibatchsize, I'm using RTX4090 (24GB) x8 with minibatch size=320 and batch size=1600, while the source code corresponds to 1024 and 5120. I'm unsure if this alone could explain the reproduction gap, or if I need to adjust the learning rate or other hyperparameters to compensate for the smaller batch size.

Then I would like to confirm a few details.

1. Is bert-144-l6-reset-custom-tokenizer fully randomly initialized from a fresh BertConfig, or does it inherit any pretrained weights in any form?
   From the released weights statistics, it looks consistent with random BERT initialization (std around 0.02), but I want to confirm this explicitly.

2. Are there any important training details not fully captured by osm_tag_embed_from_scratch.py that may materially affect the final model?
   For example:

   • exact number of GPUs used(I guess the original experiment used four?)
   • exact global batch size
   • seed / data seed
   • exact sentence-transformers / transformers versions
   • whether shard order was fixed/sorted before loading the parquet files
   • whether the published checkpoint was trained with the script exactly as released

3. Is the released BERT checkpoint bert-144-osm-tags-embed-from_scratch_torch19 expected to be reproducible from the released dataset + released init model + released training script alone, or were there additional environment/code details involved?

One thing that also confused me is that newer transformers versions complain about eval_strategy="steps" without eval_dataset, so I am not fully sure whether the released run used a slightly different code path or library version.

Any clarification would be very helpful. Thanks a lot.

[immel-f]

Hi,
I think the smaller batch size may be the main culprit, in my experiments (and in contrastive learning in general afaik) the batch size had a big impact on the training outcome. In my original trainings I used 8 GPUs with this individual batch size, I would try to get your batch size as large as possible with the CachedMultipleNegativesSymmetricRankingLoss.

The bert-144-l6-reset-custom-tokenizer is randomly initialized and has no pretrained weights, this was trickier than expected to do in sentence-transformers so I released that alongside the pretrained model. I checked in my environment which I used for the BERT training and there it used sentence-transformers ==3.1.1 and transformers==4.45.2, but I am not sure that this would have a great impact on the training outcome. The published checkpoint was trained with that script, however I might have missed some details as the focus for reproducibility was the main joint training.

[xxn967]

Thank you for your previous response and for clarifying the impact of the batch size.
I have an update regarding my reproduction efforts. To match your original settings (minibatch 1024 / batch size 5120) on my hardware, I enabled gradient checkpointing to significantly reduce VRAM usage:
model[0].auto_model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) model[0].auto_model.config.use_cache = False
With this adjustment, I was able to run the training with the exact same batch size and hyperparameters as yours. However, despite matching these critical settings and using the same library versions (sentence-transformers==3.1.1, transformers==4.45.2), I am still seeing a significant performance gap compared to your released BERT checkpoint.

I also noticed that the released model weights are labeled as ckpt-10628. I did some calculations:
Given the training set of 54,423,549 samples and 4 epochs, the total training volume is ~217.7M samples.
If using 4 GPUs with a per_device_train_batch_size of 5120 (Global Batch Size = 20,480), the total steps would be exactly 10,628.
This aligns perfectly with your checkpoint name, leading me to assume the original run was likely on 4 GPUs rather than 8. Could you please confirm if this is correct?

Since I have now aligned the batch size and environment, I am particularly curious about your comment regarding the initialization of bert-144-l6-reset-custom-tokenizer:
"This was trickier than expected to do in sentence-transformers so I released that alongside the pretrained model."
Could you please elaborate on why a standard initialization (e.g., BertModel(config) ) was not sufficient in this case? I'm wondering if there were specific nuances in how the vocabulary was handled, how the embeddings were tied, or if sentence-transformers was inadvertently pulling in some defaults that you had to manually override.

If you could share a brief snippet or explanation of how you generated that initial "reset" model, it would be extremely helpful for my understanding of the pretraining pipeline.
Thank you again for your time and for supporting the community!

[immel-f]

Your investigation regarding the number of GPUs looks correct to me, my memory might have been a little off 😅
I was starting off from an existing embedding model (I think it was this) and wanted to keep the parameters and everything around the model as close to that as possible, so I manually reset all the weights. Maybe there was an easier path to that as well :D

What is the performance gap you are experiencing with your own trained model? Is it the final trained model on AV2/nuScenes or the dummy evaluation in the pretraining of the BERT model only? If it is the latter then that might be expected as that evaluation is just a dummy test. Also note that the base model (M3TR/MapTRv2) is also quite sensitive to batch size - if you differ from the 8 x (4 GPUs) =32 batch size from the config too much that might also impact your results