Add functions to compute and combine Gaussian statistics (contribution from Samsung, AI Center, Cambridge) by rogiervd · Pull Request #2857 · speechbrain/speechbrain

rogiervd · 2025-03-11T11:19:31Z

What does this PR do?

This adds functions to compute and combines Gaussian statistics (mean and variance) on data. This comes out of me reviewing #2835, which also implements this.

However, I'm not clever enough to be confident that I fully understand this aspect of #2835, so I hereby offer an implementation which is simple enough that even I can understand it. Since the code is split up into different functions that each do one thing and can easily be tested and each require a docstring, the line count is a bit higher.

combine_gaussian_statistics and combine_gaussian_statistics_distributed compute the overall mean and variance from multiple means and variances exactly as if they were computed on all the data concatenated. After formulating the desired behaviour this way, the unit tests can test exactly this. For the DDP tests, I've copy-pasted from #2835.

@pplantinga if this looks interesting, I will still add support for std=None (assuming this is really necessary), and for masks.

Before submitting

Did you read the contributor guideline?
Did you make sure your PR does only one thing, instead of bundling different changes together?
Did you make sure to update the documentation with your changes? (if necessary)
Did you write any new necessary tests? (not for typos and docs)
Did you verify new and existing tests pass locally with your changes?
Did you list all the breaking changes introduced by this pull request?
Does your code adhere to project-specific code style and conventions?

PR review

Reviewer checklist

Is this pull request ready for review? (if not, please submit in draft mode)
Check that all items from Before submitting are resolved
Make sure the title is self-explanatory and the description concisely explains the PR
Add labels and milestones (and optionally projects) to the PR so it can be classified
Confirm that the changes adhere to compatibility requirements (e.g., Python version, platform)
Review the self-review checklist to ensure the code is ready for review

pplantinga · 2025-03-14T13:27:52Z

This code looks good, its clean and fits the speechbrain design philosophy of providing clear reference algorithms. I admit that its a bit wordy for my taste and handles edge cases that I feel are unnecessary, but perhaps that's for the best.

I'm not sure we really need to support std=None but the padding masks are a crucial part of this, so go ahead and add them.

rogiervd · 2025-03-14T13:40:11Z

Great. I've started to add support for masks, which I should be able to get in by Tuesday.

What edge cases could I get rid of you think? I do agree it has ended up wordy.

pplantinga · 2025-03-14T13:44:24Z

+        if dim == ():
+            return 1, x, torch.zeros_like(x)


Mainly I was thinking of this edge case, I'm not too concerned about someone passing an empty tuple here, and if they did I'm guessing it would break in predictable ways. But honestly this is only two lines, so, who cares?

Meaning we can probably just leave it as-is.

rogiervd

@pplantinga sorry about the delay. This should now be ready. Adding support for masks was a bit more of a hassle than I thought.

rogiervd · 2025-03-27T15:44:18Z

+        This tensor should have the same number of dimensions as "x".
+        The dimensions indicated by "dim" should have the same size as the
+        matching dimensions in "x".
+        The other dimensions should have size 1.


@pplantinga would you mind checking that this is indeed the correct specification for the mask?

Thinking about this more, especially with the number of checks we would have to do, I wonder if we shouldn't actually just pass the relative lengths here and create the mask in the correct format. This is a pretty complex mask structure that may confuse anyone trying to use this.

I think the problem is dim. I feel like writing down and checking the shape of the lengths tensor is going to be just as complicated.

rogiervd · 2025-03-27T15:45:20Z

+        If None, then statistics will be computed over all dimensions and
+        scalar-valued statistics will be returned.
+        () has the same effect as None, which is nonsensical but it consistent
+        with torch.sum and friends.


I realised that dim=() worked strangely since this is what torch.sum implements. Now I just inherit this behaviour.

Makes sense, but perhaps we don't need to mention that out in our doctsring -- I don't know if anyone will ever call this function directly (only through input norm) and if they do I doubt they would call it with dim=()

Well the function is general enough! Why restrict use to input normalisation?

pplantinga

I am of the opinion that we should re-think the way the masks are done. This also still doesn't support skipping the variance calculation.

pplantinga · 2025-03-28T13:53:01Z

+        If None, then statistics will be computed over all dimensions and
+        scalar-valued statistics will be returned.
+        () has the same effect as None, which is nonsensical but it consistent
+        with torch.sum and friends.


Makes sense, but perhaps we don't need to mention that out in our doctsring -- I don't know if anyone will ever call this function directly (only through input norm) and if they do I doubt they would call it with dim=()

pplantinga · 2025-03-28T13:55:20Z

+        This tensor should have the same number of dimensions as "x".
+        The dimensions indicated by "dim" should have the same size as the
+        matching dimensions in "x".
+        The other dimensions should have size 1.


Thinking about this more, especially with the number of checks we would have to do, I wonder if we shouldn't actually just pass the relative lengths here and create the mask in the correct format. This is a pretty complex mask structure that may confuse anyone trying to use this.

pplantinga · 2025-03-28T13:57:21Z

+    ],
+)
+@pytest.mark.parametrize("use_mask", [False, True])
+@pytest.mark.parametrize("random_seed", [20250304, 20250326, 20250327])


Do we really need to parameterize the seed? It seems like this multiplies the tests by 3 without a good cause. In general, I found it harder to debug the tests when there were too many of them because I had to scroll through so many to find the one I was debugging.

It did make a difference in the tests. Would you prefer a for-loop and an extra level of indentation? (I don't, hence the current format.)

Okay, this is fair if it makes a difference.

rogiervd · 2025-03-28T15:45:04Z

I am of the opinion that we should re-think the way the masks are done.

Isn't this how masks are done all across SpeechBrain, apart from swapping True/False maybe? For the case where dim is all dimensions except the last one.

I would personally have left dim out instead.

This also still doesn't support skipping the variance calculation.

Earlier you said

I'm not sure we really need to support std=None but the padding masks are a crucial part of this, so go ahead and add them.

so I left out support for the variance calculation. My main problem with it is that I don't know how to express in the type system that if you put in a Tensor for std or pass in compute_variance=True, you definitely get a tensor back.

pplantinga · 2025-03-28T17:50:00Z

I think the problem is dim. I feel like writing down and checking the shape of the lengths tensor is going to be just as complicated.

Perhaps you're right that the problem is dim. For the purpose of input norm (which started all of this) we have constant dim, it will always be (0, 1). The reason for me starting to make it more general was because the "sentence" norm only normed over dim=1 but that can already not be supported here due to the number being different across different features. So one option is to delete the dim parameter and always norm those two dimensions. Or we can just keep it as-is for potential future uses. This way is more complex, but its already done and we could just move ahead with it.

pplantinga

I think we should just go ahead with the mask version because it is more flexible -- handles more cases than just using the length, even if it is a little more complex.

pplantinga · 2025-03-30T21:02:43Z

+    ],
+)
+@pytest.mark.parametrize("use_mask", [False, True])
+@pytest.mark.parametrize("random_seed", [20250304, 20250326, 20250327])


Okay, this is fair if it makes a difference.

TParcollet · 2025-03-30T21:31:31Z

Feel free to merge @pplantinga

rogiervd mentioned this pull request Mar 11, 2025

Fix input normalization and global normalization variance calculation #2835

Merged

rogiervd changed the title ~~Add functions to compute and combine Gaussian statistics~~ Add functions to compute and combine Gaussian statistics (contribution from Samsung, AI Center, Cambridge) Mar 11, 2025

rogiervd force-pushed the gaussian_statistics branch from 8ecb922 to 6543aff Compare March 11, 2025 15:54

pplantinga reviewed Mar 14, 2025

View reviewed changes

Rogier van Dalen added 3 commits March 27, 2025 11:51

Add functions to compute and combine Gaussian statistics

341e8d3

Add docstring

557260a

Add support for masks

57bbb9b

rogiervd force-pushed the gaussian_statistics branch from 6543aff to 57bbb9b Compare March 27, 2025 13:37

Rogier van Dalen added 5 commits March 27, 2025 14:10

Remove the option of setting std=None

60dc766

Fix calculation with masks

87a2f3b

Test masks more properly

77aa8b8

Niceties

ea4e066

More niceties

9ee6d6a

rogiervd commented Mar 27, 2025

View reviewed changes

pplantinga reviewed Mar 28, 2025

View reviewed changes

pplantinga approved these changes Mar 30, 2025

View reviewed changes

pplantinga merged commit 9078e55 into speechbrain:develop Mar 31, 2025

Conversation

rogiervd commented Mar 11, 2025

What does this PR do?

PR review

Uh oh!

pplantinga commented Mar 14, 2025

Uh oh!

rogiervd commented Mar 14, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rogiervd left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pplantinga left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rogiervd commented Mar 28, 2025

Uh oh!

pplantinga commented Mar 28, 2025

Uh oh!

pplantinga left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

TParcollet commented Mar 30, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants