Add CategoricalMADE
#1269
Add CategoricalMADE
#1269
Conversation
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1269 +/- ##
==========================================
- Coverage 86.05% 77.87% -8.19%
==========================================
Files 118 119 +1
Lines 8672 8786 +114
==========================================
- Hits 7463 6842 -621
- Misses 1209 1944 +735
Flags with carried forward coverage won't be shown. Click here to find out more.
|
…too. log_prob has shape issues tho
…ting mixed_density estimator log_probs and sample to work as well
…rg to categorical_model
|
Hey @janfb, Currently the PR adds the As far as I can tell all functionalities of The question now is: How should I verify this works? / Which tests should I add/modify? Do you have an idea for a good toy example with several discrete variables that I could use? I have cooked up a toy simulator, for which I am getting good posteriors using SNPE, but for some reason MNLE raises a This is the simulator def toy_simulator(theta: torch.Tensor, centers: list[torch.Tensor]) -> torch.Tensor:
batch_size, n_dimensions = theta.shape
assert len(centers) == n_dimensions, "Number of center sets must match theta dimensions"
# Calculate discrete classes by assiging to the closest center
x_disc = torch.stack([
torch.argmin(torch.abs(centers[i].unsqueeze(1) - theta[:, i].unsqueeze(0)), dim=0)
for i in range(n_dimensions)
], dim=1)
closest_centers = torch.stack([centers[i][x_disc[:, i]] for i in range(n_dimensions)], dim=1)
# Add Gaussian noise to assigned class centers
std = 0.4
x_cont = closest_centers + std * torch.randn_like(closest_centers)
return torch.cat([x_cont, x_disc], dim=1)The setup: torch.random.manual_seed(0)
centers = [
torch.tensor([-0.5, 0.5]),
# torch.tensor([-1.0, 0.0, 1.0]),
]
prior = BoxUniform(low=torch.tensor([-2.0]*len(centers)), high=torch.tensor([2.0]*len(centers)))
theta = prior.sample((20000,))
x = toy_simulator(theta, centers)
theta_o = prior.sample((1,))
x_o = toy_simulator(theta_o, centers)NPE: trainer = SNPE()
estimator = trainer.append_simulations(theta=theta, x=x).train(training_batch_size=1000)
snpe_posterior = trainer.build_posterior(prior=prior)
posterior_samples = snpe_posterior.sample((2000,), x=x_o)
pairplot(posterior_samples, limits=[[-2, 2], [-2, 2]], figsize=(5, 5), points=theta_o)and the equivalent MNLE: trainer = MNLE()
estimator = trainer.append_simulations(theta=theta, x=x).train(training_batch_size=1000)
mnle_posterior = trainer.build_posterior(prior=prior)
mnle_samples = mnle_posterior.sample((10000,), x=x_o)
pairplot(mnle_samples, limits=[[-2, 2], [-2, 2]], figsize=(5, 5), points=theta_o)Hoping this makes sense. Lemme know if you need clarifications anywhere. Thanks for your feedback. |
|
Hey @janfb, |
There was a problem hiding this comment.
thanks a lot for tackling this @jnsbck! 👏
Please find below some comments and questions.
There might be some misunderstanding about variables and categories on my side. We can have a call if that's more efficient than commenting here.
| epsilon=1e-2, | ||
| custom_initialization=True, | ||
| embedding_net: Optional[nn.Module] = nn.Identity(), | ||
| ): |
There was a problem hiding this comment.
please add docstring with brief explanation and arg list.
| categories, # Tensor[int] | ||
| hidden_features, | ||
| context_features=None, | ||
| num_blocks=2, | ||
| use_residual_blocks=True, | ||
| random_mask=False, | ||
| activation=F.relu, | ||
| dropout_probability=0.0, | ||
| use_batch_norm=False, | ||
| epsilon=1e-2, | ||
| custom_initialization=True, |
| if custom_initialization: | ||
| self._initialize() | ||
|
|
||
| def forward(self, inputs, context=None): |
| embedded_context = self.embedding_net.forward(context) | ||
| return super().forward(inputs, context=embedded_context) | ||
|
|
||
| def compute_probs(self, outputs): |
There was a problem hiding this comment.
add types and short docstring.
| # outputs (batch_size, num_variables, num_categories) | ||
| def log_prob(self, inputs, context=None): |
There was a problem hiding this comment.
remove inline comment and instead add types, return types and docstring with details on the return dimensions if needed.
| self.num_variables = len(categories) | ||
| self.num_categories = int(max(categories)) | ||
| self.categories = categories |
There was a problem hiding this comment.
I am confused by the notion of these three variables.
what's the difference between variables and categories here?
There was a problem hiding this comment.
what type does categories have?
| "this is not the case for the passed `x` do not use this function.", | ||
| stacklevel=2, | ||
| ) | ||
| # Separate continuous and discrete data. | ||
| cont_x, disc_x = _separate_input(batch_x) | ||
| num_disc = int(torch.sum(_is_discrete(batch_x))) |
There was a problem hiding this comment.
please add num_categorical_columns as input arg to this function to enable users to pass this number. Inferring it from batch_x should be the fallback with warning.
| z_score_y="none", # y-embedding net already z-scores. | ||
| num_hidden=hidden_features, | ||
| num_layers=hidden_layers, | ||
| embedding_net=embedding_net, |
There was a problem hiding this comment.
pass num_categories here as well.
| elif categorical_model == "mlp": | ||
| assert num_disc == 1, "MLP only supports 1D input." | ||
| discrete_net = build_categoricalmassestimator( | ||
| disc_x, | ||
| batch_y, | ||
| z_score_x="none", # discrete data should not be z-scored. | ||
| z_score_y="none", # y-embedding net already z-scores. | ||
| num_hidden=hidden_features, | ||
| num_layers=hidden_layers, | ||
| embedding_net=embedding_net, | ||
| ) |
There was a problem hiding this comment.
more generally, isn't the MLP a special case of the MADE? can't we absorb them into one class?
There was a problem hiding this comment.
this should eventually be integrated with the MNLE tutorial in 12_iid_data_and_permutation_invariant_embeddings.ipynb
|
Cool, thanks for all the feedback! A quick call would be great, also to discuss suitable tests for this. Will reach out via email and tackle the straight forward things in the meantime. |
What does this implement/fix? Explain your changes
This implements a
CategoricalMADEto generelize MNLE to multiple discrete dimensions addressing #1112.Essentially adapts
nflows's MixtureofGaussiansMADE to autoregressively model categorical distributions.Does this close any currently open issues?
Fixes #1112
Comments
I have already discussed this with @michaeldeistler.
Checklist
Put an
xin the boxes that apply. You can also fill these out after creatingthe PR. If you're unsure about any of them, don't hesitate to ask. We're here to
help! This is simply a reminder of what we are going to look for before merging
your code.
guidelines
with
pytest.mark.slow.guidelines
main(or there are no conflicts withmain)