concept bottleneck generative models - #71
Open
gdefe wants to merge 1 commit into
Open
Conversation
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
There was a problem hiding this comment.
Pull request overview
Adds a high-level Concept Bottleneck Generative Model (CBGM, VAE-style) to the library, reusing the existing CEM mixture machinery while extending the mid-level Pyro inference utilities to better support matrix-shaped PGM variables and deterministic (Delta) nodes.
Changes:
- Introduces
ConceptBottleneckGenerativeModel(CBGM) and example scripts demonstrating conditional generation/steering on ColorMNIST. - Refactors/extends concept–embedding mixing layers and adds the CBGM orthogonality penalty (
concept_orthogonality). - Updates Pyro inference utilities to handle matrix-shaped event parameters and Delta parameter extraction.
Reviewed changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated 8 comments.
Show a summary per file
| File | Description |
|---|---|
| torch_concepts/nn/modules/mid/inference/pyro/utils.py | Adds Delta parameter extraction for deterministic sites. |
| torch_concepts/nn/modules/mid/inference/pyro/base.py | Flattens matrix-shaped event parameters to avoid Pyro plate/batch collisions. |
| torch_concepts/nn/modules/low/predictors/mix.py | Introduces shared mixing base + adds CBGM bottleneck layer. |
| torch_concepts/nn/modules/high/models/cem.py | Renames reordered concept axis variable for clarity. |
| torch_concepts/nn/modules/high/models/cbgm.py | Adds the CBGM high-level model and its variational guide wiring. |
| torch_concepts/nn/functional.py | Adds concept_orthogonality penalty used by CBGM. |
| torch_concepts/nn/init.py | Exports the new CBGM model/layer. |
| examples/utilization/2_model/15_concept_bottleneck_generative_model.py | End-to-end CBGM training demo on ColorMNIST. |
| examples/utilization/0_layer/6_cbgm_layer_equivalence.py | Low-level equivalence check against the reference CBGM implementation. |
| doc/modules/nn.functional.rst | Documents concept_orthogonality. |
| doc/modules/low_level_api.rst | Documents MixConceptEmbeddingToEmbedding. |
| doc/modules/high_level_api.rst | Documents ConceptBottleneckGenerativeModel. |
Comments suppressed due to low confidence (6)
torch_concepts/nn/modules/low/predictors/mix.py:235
MixConceptEmbeddingToEmbeddingclaims to output the full CBGM bottleneckw = [w_1, …, w_k, w_unk](and the doctest calls it withunknown=...), butforward()currently returns only the mixed supervised contexts (krows) and does not incorporate the unsupervised context. This makes the doc/examples incorrect and prevents reproducing the reference CBGM bottleneck.
def forward(
self,
concepts: torch.Tensor,
embeddings: torch.Tensor,
) -> torch.Tensor:
torch_concepts/nn/modules/high/models/cbgm.py:332
MixConceptEmbeddingToEmbeddingis used here to build the bottleneck, butunknownis not included as a parent or aggregated into the layer'sembeddingsinput. With the CBGM bottleneck definition, the mixer needs the unsupervised context so it can returnk+1embeddings (and socontextcan be a single deterministic node).
mixing_cpd = ParametricCPD(
variable=mixing,
parents=[*concepts, *embeddings],
parametrization={"value": MixConceptEmbeddingToEmbedding(
in_concepts=reordered_axis, # require Annotations as in_concepts
in_embeddings=self.embedding_size,
)},
aggregate=mix_parents,
)
torch_concepts/nn/modules/high/models/cbgm.py:346
- The decoder CPD currently takes
[mixing, unknown]and concatenates them manually. Once the bottleneck is represented as a single deterministiccontextvariable, the decoder should depend oncontextalone; otherwise the graph keeps two parallel representations of the bottleneck and examples that useout.value['context']will fail.
decoder_cpd = ParametricCPD(
variable=observed,
parents=[mixing, unknown],
parametrization=self._flexible_parametrization(
variable=observed,
first=pyc.nn.Sequential(nn.Flatten(start_dim=-2), self.decoder),
second='auto',
),
aggregate=cat_embeddings,
)
torch_concepts/nn/modules/high/models/cbgm.py:351
- The BayesianNetwork assembly still references
mixingandmixing_cpd. After switching to a singlecontextnode andcontext_cpd, the variables/factors lists need to be updated; otherwise the model will reference undefined names and the context node won't be part of the graph/output.
return BayesianNetwork(
variables=[latent, *embeddings, unknown, *concepts, mixing, observed],
factors=[latent_cpd, *emb_encoders, *c_encoders, mixing_cpd, decoder_cpd],
)
examples/utilization/2_model/15_concept_bottleneck_generative_model.py:125
- With
observation=Bernoulli, the reconstruction term should be computed fromout.probs['input'](orout.logits['input']) using a Bernoulli-appropriate loss (BCE). Usingout.loc['input']will fail for Bernoulli outputs and also doesn't match the stated observation model.
# Reconstruction and the Gaussian KL of the guide vs the N(0, I) prior.
recon = F.mse_loss(
out.loc["input"], x.flatten(1), reduction="none"
).sum(-1).mean()
examples/utilization/2_model/15_concept_bottleneck_generative_model.py:105
- When using
observation=Bernoullithe decoder must output valid probabilities in[0, 1]because the model is configured to parameterize discrete variables viaprobs. The example decoder currently has noSigmoid, so it will emit unconstrained reals that are invalid Bernoulli parameters.
decoder=torch.nn.Sequential(
MLP(context_size, 64, math.prod(dataset.n_features), n_layers=2, activation="leaky_relu"),
torch.nn.Unflatten(-1, dataset.n_features)
)
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| self.binary_mask = torch.from_numpy(np.array(in_concepts.types) != 'continuous') | ||
| cumsum = torch.cumsum(self.cardinalities_expanded, dim=0) | ||
| start_positions = cumsum - self.cardinalities_expanded | ||
| bernoulli_mask = self.cardinalities_expanded == 1 & self.binary_mask |
Comment on lines
+78
to
+80
| observation : type, default ``torch.distributions.Bernoulli`` | ||
| Distribution family of the generated variable. ``Bernoulli`` for images | ||
| in ``[0, 1]``; ``Normal`` adds a learned scale head. |
Comment on lines
+255
to
+269
| # The pre-defined concepts are incomplete in a generative setting, so the | ||
| # bottleneck carries one extra, unsupervised context embedding. | ||
| unknown = EmbeddingVariable( | ||
| "unknown", | ||
| distribution=Delta, | ||
| shape=(1, self.embedding_size), | ||
| ) | ||
| ordered_names = [m for cvar in concepts for m in cvar.members] | ||
| reordered_axis = self.concept_annotations.subset(ordered_names) | ||
| n_concepts = len(ordered_names) | ||
| mixing = EmbeddingVariable( | ||
| "mixing", | ||
| distribution=Delta, | ||
| shape=(n_concepts, self.embedding_size), | ||
| ) |
Comment on lines
+91
to
+106
| model = ConceptBottleneckGenerativeModel( | ||
| input_size=dataset.n_features, | ||
| annotations=dataset.annotations, | ||
| backbone=backbone, | ||
| encoder=MLP( | ||
| backbone.out_features, | ||
| LATENT_SIZE, | ||
| n_layers=1 | ||
| ), | ||
| latent_size=LATENT_SIZE, | ||
| embedding_size=EMBEDDING_SIZE, | ||
| decoder=torch.nn.Sequential( | ||
| MLP(context_size, 64, math.prod(dataset.n_features), n_layers=2, activation="leaky_relu"), | ||
| torch.nn.Unflatten(-1, dataset.n_features) | ||
| ) | ||
| ) |
Comment on lines
+136
to
+142
| pyc_bottleneck = layer( | ||
| concepts=torch.cat(probs, dim=-1), | ||
| embeddings=torch.cat( | ||
| [c.unflatten(-1, (-1, EMB_SIZE)) for c in contexts[:-1]], dim=-2 | ||
| ), | ||
| unknown=contexts[-1].unflatten(-1, (1, EMB_SIZE)), | ||
| ) |
Comment on lines
+167
to
+170
| class MixConceptEmbeddingToEmbedding(MixConceptEmbedding): | ||
| """Concept bottleneck layer of a Concept Bottleneck Generative Model. | ||
|
|
||
| Mixes each concept's state embeddings by its predicted probability — the |
Comment on lines
+182
to
+186
| def concept_orthogonality(context: torch.Tensor, n_concepts: int) -> torch.Tensor: | ||
| """ | ||
| Orthogonality penalty between the concept contexts and the unsupervised one. | ||
|
|
||
| The concept bottleneck of a generative model ends in an *unsupervised* |
Comment on lines
+201
to
+206
| >>> concepts = torch.rand(4, 12) # (batch, 10 + 2 state scores) | ||
| >>> embeddings = torch.randn(4, 12, 16) # (batch, states, m) | ||
| >>> unknown = torch.randn(4, 1, 16) # (batch, 1, m) | ||
| >>> | ||
| >>> layer(concepts=concepts, embeddings=embeddings, unknown=unknown).shape | ||
| torch.Size([4, 3, 16]) |
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.
Validation: