Skip to content

Soft contrastive loss, scalable alignment training implementation #118

Open
gabrieletijunaityte wants to merge 173 commits into
developfrom
train/alignment
Open

Soft contrastive loss, scalable alignment training implementation #118
gabrieletijunaityte wants to merge 173 commits into
developfrom
train/alignment

Conversation

@gabrieletijunaityte

@gabrieletijunaityte gabrieletijunaityte commented Jun 16, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Sorry for a big, complex PR. Alignment training needed quite a few fixes.

Changes:

  1. Validation loss is now calculated with n randomly sampled location templates averaged instead of all. Allows for much faster training, with just a bit less stable validation loss curves. Implemented:
    • Add n_captions_for_validation parameter to caption builder.
    • Random location caption generation for n_captions_for_validation randomly sampled captions in Captionuilder.sample_multiple_or_all()
  2. Some minor fixes for adopted geo encoders (e.g., adoption flag).
  3. Datamodule changes:
    • Add _setup_flag parameter to only setup once.
    • Datamodule can return name_loc in the get_item now → used for debugging.
    • Move data splitting into setup part (allows to first filter for any missing data, and only then split the remaining data).
    • Add data size per split print statements.
  4. CNN from SSL4EO fixes.
  5. Fix LLM implementation.
  6. Text alignment model now has geo_encoder (which can be adopted from the prediction model) and text_encoder as before, with additional new geo and text adapters.

New features:

  1. Notebook that computes aux column statistics for standardization. It uses only training locations. (Included these aux files here).
  2. Notebook that computes average embeddings for tessera and aef tiles and saves them in a CSV file.
  3. Notebook on inspecting text encoder embeddings.
  4. New supported EO modalities in geo-encoders: avr embeddings from a CSV file.
  5. Implement soft contrastive loss, which allows environmentally similar locations to be treated as soft positives. Implemented necessary changes:
    • Caption builder now has an aux column standardization file parameter. Standardization happens in the collate_fn (allows batch-wise computation).
    • Caption builder can be requested to return aux column Idx that have been used in the template.
  6. Unlabelled data implementation from @vdplasthijs
  7. Move MLP projector from base_geo_encoder to a new class of projectors_adapters.

Fixes #<issue_number>

Before submitting

  • Did you make sure title is self-explanatory and the description concisely explains the PR?
  • Did you make sure your PR does only one thing, instead of bundling different changes together?
  • Did you list all the breaking changes introduced by this pull request?
  • Did you test your PR locally with pytest command?

@gabrieletijunaityte gabrieletijunaityte marked this pull request as ready for review June 17, 2026 10:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants