Enforce consistent policy for parameter freezing - #456
Open
nkbranigan wants to merge 1 commit into
Open
Conversation
- Fitting some models with EM silently ignores `trainable=False`. Further, the models that do respect the flag disagree on the policy: for example, the GMM HMMs reject freezing all emission parameters while `GaussianHMM` allows it.
- New policy:
- When using EM with a closed-form M-step: HMMs can freeze all or no emission parameters (and independently freeze the initial and transition parameters in any combination — no change to this behavior), and LGSSMs can freeze all or no parameters
- Unsupported freezes raise, with instruction to use `fit_sgd`, which supports any combination of frozen parameters
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.
Summary
trainable=False. Further, the models that do respect the flag disagree on the policy: for example, the GMM HMMs reject freezing all emission parameters whileGaussianHMMallows it.Details
fit_sgd, which supports any combination of frozen parameterstrainableis static), so EM's XLA programs are identical before and after this PR when all parameters are trainable (verified by comparingjaxprfor all 9 affected models)