fix(architectures): apply the Mamba block in MambaTab.forward - #435
Open
ChrisW09 wants to merge 2 commits into
Open
fix(architectures): apply the Mamba block in MambaTab.forward#435ChrisW09 wants to merge 2 commits into
ChrisW09 wants to merge 2 commits into
Conversation
self.mamba was constructed in __init__ but forward() went initial_layer -> unsqueeze -> norm_f -> activation -> squeeze -> head and never invoked it. Every MambaTab classifier/regressor/LSS silently trained as a single linear layer with an MLP head, while all Mamba parameters sat in the optimizer receiving no gradients. Fixes #412 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
model._task_model and model._data_module are Optional-typed internals. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #412
Problem
MambaTab.__init__buildsself.mamba(eitherMambaorMambaOriginaldepending onconfig.mamba_version), butforward()never calls it:So every
MambaTabClassifier/MambaTabRegressor/MambaTabLSStrains as a linear layer + LayerNorm + activation + MLP head. The Mamba parameters are registered, consume memory and optimizer state, and receive zero gradients, while the model silently underperforms what it claims to be. (OnlyBaseModel.encode()for pretraining ever touchedself.mamba.)Fix
Apply
self.mambato the sequence-shaped tensor, between the activation and the squeeze back to(B, d_model).Tests
New
tests/test_mambatab_backbone.pyfits a model, runs one manual forward/backward, and asserts every.mamba.parameter has a gradient. This shape of test would also have caught the SAINT dead-LayerNorm issue (#416) — worth generalizing across architectures in a follow-up.🤖 Generated with Claude Code