Repository navigation
Conversation
The "random-signs" entry of the scaling-init table built a new tensor with torch.sign(torch.randn_like(x)) and discarded it, so the r/s scaling parameters of LinearBatchEnsembleLayer, RNNBatchEnsembleLayer and MultiHeadAttentionBatchEnsemble kept the uninitialized torch.empty memory (often zeros, which blocks the gradient of the shared weight matrix). Add an in-place _random_signs_ initializer that fills the tensor with -1/+1 values, matching the other torch.nn.init entries. Fixes #497
TabMConfig and the batch-ensemble layers now say what 'ones', 'random-signs' and 'normal' initialize the scaling factors to, and the TabM page names the values. The TabM test checks |r| == 1 instead of expecting both signs in a 4 x 3 tensor, which can draw a single sign. Refs #497
This branch has not been deployed
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
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 #497
Problem
The
"random-signs"option ofscaling_initwas implemented aslambda x: torch.sign(torch.randn_like(x)). That lambda returns a new tensor, whichreset_parametersthrows away. The batch-ensemble scaling parametersrandstherefore kept the uninitializedtorch.emptymemory. That memory was all zeros in my runs, so the first TabM layer only output its bias and its shared weight matrix got no gradient. The same code is inLinearBatchEnsembleLayer,RNNBatchEnsembleLayerandMultiHeadAttentionBatchEnsemble.Fix
A small in-place initializer,
_random_signs_(tensor)indeeptab/nn/blocks/common.py, fills the tensor with-1/+1undertorch.no_grad(). Like thetorch.nn.initfunctions it modifies its argument and returns it. The threescaling_init_fntables use it in place of the lambda."ones"and"normal"are unchanged.Tests
tests/nn/blocks/test_common_blocks.py::TestRandomSignsScalingInit: for all three layers,r/scontain only-1and+1. The test makestorch.emptyreturn NaN-filled memory, so leftover zeros or uninitialized values can't pass by accident.tests/architectures/test_architecture_contracts.py::test_tabm_random_signs_scaling_init_trains_first_layer: a TabM built withscaling_init="random-signs"has±1scaling factors, and the first layer's weight matrix gets a non-zero gradient.All of these fail on
mainand pass with this change.Docs
TabMConfig.scaling_initand theLinearBatchEnsembleLayerandRNNBatchEnsembleLayerdocstrings describe the three options:'ones'sets the scaling factors to 1,'random-signs'draws each from {-1, +1}, and'normal'draws them from N(0, 1).MultiHeadAttentionBatchEnsemblealready did.LinearBatchEnsembleLayer'sscaling_initannotation now includes"normal", which it already accepted.docs/model_zoo/stable/tabm.mdnames the three initializations instead of "Diversity initialization for scaling factors".Checks
just test: 3240 passed, 46 skipped, 28 xfailed; coverage 91.29%just docs: build succeeded (warnings treated as errors)just check: all hooks passed (ruff, prettier, pyright)