feat!(fastsiam): more faithful FastSiam implementation

[giuliano-macedo]

Checklist before requesting a review

• I have performed a self-review of my code.
• I have added tests to this functionality.
• All tests are passing.
• I have updated the documentation as needed.
• I have followed the project's coding guidelines.

---

Hi, first of all thanks for the implementation.

Motivation

While reading the FastSiam implementation, I've noticed some divergences when comparing it with the original paper.

Although I have seem good results with the current implementation, I still think it is different from the original paper authors intent.

So this PR is more of a faithful implementation of the paper, than a de-facto "fix".

The main divergences that I observed are:

• 1. It is not a Siamese network: The current implementation does not share weights between the branches.
     The paper explicitly states:

FastSiam extends upon SimSiam and addresses a core weakness of it, namely target instability.

Since SimSiam is an implementation that shares weights between branches, this doesn't seem to be the same algorithm described in the paper.

• 2. $K=1$ does not equate to SimSiam: Since the code uses a target network updated via momentum, if we set $k=1$, the network does not equate to SimSiam (unlike the mathematical formulation in the paper).
     In fact, this implementation seems to decay into the BYOL architecture.

And also, a minor issue that doesn't really impact these points is that the update_target_branch method was defined, but not called anywhere.
If the user simply just pass the model to trainer.fit(), it would equate to dead code and the target network will never be updated.

Summary of the changes

For $k=3$ I've inlined the loss function because that will probaly be the most common value for $k$ that the user will take, therefore inlining it would be faster than the fallback for arbitrary $k$ that I
used Lightly's implementation as a base.

I've made some breaking changes to the arguments of this class because some of them weren't used (like test_metric and num_classes), others were EMA-related (momentum), and I've changed the case of the K
argument to follow the style-guide of other classes in Minerva that uses lower case instead.

Also I've the flatten and avg_pooling: parameter that is useful when not using a CNN-based backend, a similar setup that the SimCLR class uses.

And finally, I've added a simple check to garantee the user passes $k+1$ views in his Dataset, because this algorithm assumes that the user will use k views for the target branch, and 1 view for the prediction branch, therefore $k+1$.

Reference total loss derived from the paper

To help with the review, here are the total loss for specific $k$ that I derived from the paper, and the general formula:

$\mathcal{D}=$ negative cosine similarity
$\text{sg}(x)=$ stop gradient

$K$	Formula
1	$\tfrac{1}{2}\left[\mathcal{D}(p_1,\text{sg}(z_2))+\mathcal{D}(p_2,\text{sg}(z_1))\right]$ (same as SimSiam loss)
2	$\tfrac{1}{3}\left[\mathcal{D}!\left(p_1,\tfrac{\text{sg}(z_2)+\text{sg}(z_3)}{2}\right)+\mathcal{D}!\left(p_2,\tfrac{\text{sg}(z_1)+\text{sg}(z_3)}{2}\right)+\mathcal{D}!\left(p_3,\tfrac{\text{sg}(z_1)+\text{sg}(z_2)}{2}\right)\right]$
3	$\tfrac{1}{4}\left[\mathcal{D}!\left(p_1,\tfrac{\text{sg}(z_2)+\text{sg}(z_3)+\text{sg}(z_4)}{3}\right)+\mathcal{D}!\left(p_2,\tfrac{\text{sg}(z_1)+\text{sg}(z_3)+\text{sg}(z_4)}{3}\right)+\mathcal{D}!\left(p_3,\tfrac{\text{sg}(z_1)+\text{sg}(z_2)+\text{sg}(z_4)}{3}\right)+\mathcal{D}!\left(p_4,\tfrac{\text{sg}(z_1)+\text{sg}(z_2)+\text{sg}(z_3)}{3}\right)\right]$
General	$\tfrac{1}{K+1}\sum_{i=1}^{K+1} \mathcal{D}\left(p_i,\tfrac{1}{K}\sum_{j=1, j\neq i}^{K+1}\text{sg}(z_j)\right)$

Final Notes

Please let me know if my assertions are incorrect, also, if there is any specific context behind the original design choices that I might have missed, I am more than happy to discuss this further and make any necessary adjustments based on your feedback!

[coderabbitai]

📝 Walkthrough

Walkthrough

FastSiam is refactored from a two-branch momentum-target architecture to a single-branch projector/predictor pipeline. The module now defines projection_head and prediction_head, and forward returns (z.detach(), p). Loss is computed in _single_step via cosine-similarity across k+1 multi-view inputs. Tests are updated to match the new interface.

Changes

FastSiam Single-Branch Refactor

Layer / File(s)	Summary
Constructor and forward pass redesign minerva/models/ssl/fastsiam.py	__init__ replaces prediction/target branch members with projection_head and prediction_head; adds flatten, avg_pooling, k, lr, loss_fn parameters; removes momentum, test_metric, num_classes. forward becomes a single-view pipeline returning (z.detach(), p). Docstring updated accordingly.
Loss computation and training/val/test steps minerva/models/ssl/fastsiam.py	_single_step_arbitrary_k stacks per-view outputs and aggregates cosine-similarity loss using average-of-other-views targets. _single_step validates len(batch) == self.k + 1, uses an optimized inline path for k==3, otherwise delegates to _single_step_arbitrary_k. training_step, validation_step, test_step call _single_step with a log prefix.
Updated and new tests tests/models/ssl/test_fastsiam.py	Test setup updated for k=2. test_forward passes a single view and asserts (z, p) shapes. New/rewritten tests: test_single_step_arbitrary_k, test_training_step, test_unexpected_k_error (wrong view count raises RuntimeError), test_single_step_k_equals_3, test_avg_pooling_flag, test_flatten_flag.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title accurately summarizes the main change: a more faithful FastSiam implementation.
Description check	✅ Passed	The PR description is mostly complete, with motivation, change summary, checklist, and final notes covering the template’s main requirements.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 7

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@minerva/models/ssl/fastsiam.py`:
- Around line 105-106: The `avg_pooling` docstring in `FastSiam` does not match
the behavior in `forward`, where pooling happens before flattening and is not
conditional on `flatten=True`. Update the parameter description near
`avg_pooling` to reflect the actual control flow in `FastSiam.forward`, making
it clear that average pooling is applied independently of the flatten option.
- Around line 176-180: The loss normalization in the FastSiam loss computation
is incorrect because the accumulated terms from the loop are being divided by
self.k and then shifted by 1, which mis-scales losses for arbitrary k. Update
the return in the loss calculation inside the FastSiam method that builds loss
from ps and zs so it averages over all k+1 accumulated terms instead of applying
a constant offset, using the existing loss accumulation logic and the
self.k/self.k + 1 loop bounds as the guide.
- Around line 135-143: The SimSiam projector in the FastSiam model is still
hard-coded to use 512 instead of the constructor’s hid_dim, so only the
prediction head respects user configuration. Update the projection_head
construction in FastSiam to use hid_dim as the hidden layer size, matching how
prediction_head already wires hid_dim through, so both heads follow the
configured dimensions consistently.

In `@tests/models/ssl/test_fastsiam.py`:
- Around line 26-35: The current FastSiam forward test only checks output
shapes, so it can miss regressions where the stop-gradient behavior is removed.
Update test_forward in the FastSiam test to also verify the contract of
FastSiam.forward by asserting the returned z is detached (for example, that it
does not require gradients) while p still remains part of autograd. Use the
existing self.model(view) call and the z, p outputs to place these assertions
alongside the current shape checks.
- Around line 120-140: The test for FastSiam flatten behavior is not actually
validating the flatten=False path because the mock backbone already returns a 2D
tensor, so the model can still pass even if it always flattens. Update
test_flatten_flag in test_fastsiam.py to use a non-flattened backbone output
shape from mock_backbone (for example a 3D tensor with last dimension 2048),
then assert that FastSiam preserves that extra dimension when flatten is
disabled while still producing the expected z and p outputs. Keep the check
focused on the flatten attribute and the forward path in FastSiam.
- Around line 104-119: The FastSiam pooling test currently uses a backbone
output with 1x1 spatial size, so the avg_pooling=False path is not actually
exercised. Update test_avg_pooling_flag in test_fastsiam.py to make
self.mock_backbone return a larger spatial feature map and adjust in_dim
accordingly, then keep the assertions on model.global_avg_pool and output shapes
so the forward path without pooling is truly validated.
- Around line 37-48: The current loss tests only assert that
`_single_step_arbitrary_k` and the `k == 3` path return a scalar tensor, so they
can miss arithmetic regressions like an incorrect extra offset or scaling.
Update `test_single_step_arbitrary_k` and the corresponding
`test_single_step_k3` (the other loss-path test) to use a numeric oracle by
either computing the expected loss manually or stubbing `forward`/`loss_fn`
deterministically, then assert the returned value matches that expected number
for both paths.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 182d823d-61ca-4572-84f1-ee2511e4fddc

📥 Commits

Reviewing files that changed from the base of the PR and between f23918a and 3822911.

📒 Files selected for processing (2)

• minerva/models/ssl/fastsiam.py
• tests/models/ssl/test_fastsiam.py