conditional gfn #188

josephdviviano · 2024-09-25T15:11:25Z

Supports conditioning on a tensor of shape=[n_trajectories, n_cond_dims]. This is passed by the user during a call to the sampler.

Implemented for all GFlowNets. Note that the current version expects a particular kind of estimator. I can imagine this will lead to future changes - e.g., we should have some Estimators which expect huggingface models, so we can use them to produce conditioning vectors / to initialize the policy (this will obviously be a future PR).

Note that the conditioning is useless in my example, we should have a better use-case envisioned for the demo. The demo currently is not complete for all GFlowNet types.

…ionally contains a tensor of conditioning vectors (one per trajectory)

…itioning into PB and PF computation

…ule can now accept raw tensors

…bute of the trajectory

josephdviviano · 2024-09-25T15:32:54Z

Don't worry about the tests - they should be easy to fix.

I can make the chances for DB, Sub-TB, and FM pretty easily if we agree this is a good approach, before a proper review.

saleml · 2024-09-25T17:50:21Z

src/gfn/modules.py

+
+    or
+
+    $s \mapsto (P_B(s' \mid s, c))_{s' \in Parents(s)}$.


might be worth mentioning that this is a s very specific conditioning use-case, where the condition is encoded separately, and embeddings are concatenated.

I don't think we can do a generic one, but this should be enough as an example !

What other conditioning approaches would be worth including? Cross attention?

In general I would think the conditioning should be embedded / encoded separately --- or would the conditioning just need to be concatenated to the state before input? I could add support for that.

I don't think there is an exhaustive list of ways we can process the condition. What you have is great as an example. I suggest you just add a comment or doc that the user might want to write their own module

saleml · 2024-09-25T17:51:04Z

src/gfn/samplers.py

@@ -68,7 +67,28 @@ def sample_actions(
                the sampled actions under the probability distribution of the given
                states.
        """
-        estimator_output = self.estimator(states)
+        # TODO: Should estimators instead ignore None for the conditioning vector?


wouldn't it be cleaner with fewer if else blocks ?

Yes there's a bit of cruft with all the if-else blocks, but as it stands an estimator can either accept one or two arguments and I think it's good if it fails noisily... what do you think?

Ok ! makes sense.

I added these exception_handlers to reduce the cruft.

saleml · 2024-09-25T20:04:28Z

LGTM! Looking forward to test this feature

…g conditioning

…tePolicyEstimator

saleml

I'm happy to see this being added to the library. Great work! Great code design, and thanks for factorizing a few other things, including the context managers / error handlers.

I left a few comments and a suggestion for the script.

saleml · 2024-10-05T12:29:51Z

src/gfn/gflownet/base.py

@@ -32,41 +35,41 @@ class GFlowNet(ABC, nn.Module, Generic[TrainingSampleType]):
    def sample_trajectories(
        self,
        env: Env,
-        n_samples: int,


it looks like you're handling the conditioning input to this function as a kwarg, whereas sampler's sample_trajectories have an explicit conditioning input. I'm wondering if you have a particular reason for this choice

I think maybe all functions should use an explicit conditioning kwarg, what do you think? I can make those changes.

I agree that it would be cleaner

it should be done now, let me know if i missed something.

saleml · 2024-10-05T12:38:41Z

tutorials/examples/train_conditional.py

+        conditioning = torch.rand((batch_size, 1))
+        conditioning = (conditioning > 0.5).to(torch.float)  # Randomly 1 and zero.
+
+        trajectories = gflownet.sample_trajectories(


pylance is not happy with these two variables. I'm wondering if this is due to using **kwargs (see my comment in the base.py file). If so, it would be nice to decide whether we should not care hereafter about pylance and co, or care, in which case, no more kwargs in the code base

so I think I fixed this but moving to **kwargs: Any but we have a multitude of other harder to handle pylance issues that I'm not sure what to do about and warrants a discussion bigger than the scope of this PR, I think.

saleml · 2024-10-05T12:48:41Z

tutorials/examples/train_conditional.py

+    print("+ Training Conditional {}!".format(type(gflownet)))
+    for i in (pbar := tqdm(range(n_iterations))):
+        conditioning = torch.rand((batch_size, 1))
+        conditioning = (conditioning > 0.5).to(torch.float)  # Randomly 1 and zero.


IIUC, the conditioning doesn't change anything in this example. While this file is a great way to show how one can code their conditional gflownet, what do you think of slightly altering the setting here, e.g. by making the environment conditional (e.g. hide one of the 4 modes if conditioning=1), and then, post-training, have some validation where we compare the resulting pair of distributions to the pair of target distributions ?

yes, you're right. can we save this for a follow up PR?

I filed the issue here. if you agree I'd like to do this work separately.

#190

saleml · 2024-10-05T12:49:10Z

tutorials/examples/train_conditional.py

+    gflownet = build_tb_gflownet(environment)
+    train(environment, gflownet)
+
+    gflownet = build_db_gflownet(environment)
+    train(environment, gflownet)
+
+    gflownet = build_db_mod_gflownet(environment)
+    train(environment, gflownet)
+
+    gflownet = build_subTB_gflownet(environment)
+    train(environment, gflownet)
+
+    gflownet = build_fm_gflownet(environment)
+    train(environment, gflownet)


argparse this?

saleml · 2024-10-05T12:50:10Z

src/gfn/gflownet/trajectory_balance.py

I'm pleasantly surprised no change is needed for the LogPartitionVarianceLoss. Right?

I don't think we need the conditioning information here, and I agree it's nice that the code naturally reflected that. Please correct me if I misunderstand this loss.

saleml · 2024-10-17T05:08:57Z

src/gfn/gflownet/flow_matching.py

-    ) -> tuple[DiscreteStates, DiscreteStates, torch.Tensor]:
+    def to_training_samples(self, trajectories: Trajectories) -> Union[
+        Tuple[DiscreteStates, DiscreteStates, torch.Tensor, torch.Tensor],
+        Tuple[DiscreteStates, DiscreteStates, None, None],


saleml · 2024-10-17T05:10:06Z

src/gfn/modules.py

@@ -240,13 +240,20 @@ def __init__(
        self.conditioning_module = conditioning_module
        self.final_module = final_module

-    def forward(
-        self, states: States, conditioning: torch.tensor
+    def _forward_trunk(


is what you call trunk the same thing I called torso before ?

Yeah -- let me unify the naming

saleml · 2024-10-17T05:10:55Z

LGTM!

Thanks for the PR

josephdviviano added 9 commits September 25, 2024 10:56

example of conditional GFN computation with TB only (for now)

6e8dc4d

should be no change

39fb5ee

Trajectories objects now have an optional .conditonal field which opt…

2bc2263

…ionally contains a tensor of conditioning vectors (one per trajectory)

small changes to logz paramater handling, optionally incorporate cond…

99afaf3

…itioning into PB and PF computation

logZ is optionally computed using a conditioning vector

e6d25a0

NeuralNets now have input/output dims

2c72bf9

added a ConditionalDiscretePolicyEstimator, and the forward of GFNMod…

580c455

…ule can now accept raw tensors

added conditioning to sampler, which will save the tensor as an attri…

a74872f

…bute of the trajectory

black

056d935

josephdviviano added the enhancement New feature or request label Sep 25, 2024

josephdviviano requested a review from saleml September 25, 2024 15:11

josephdviviano self-assigned this Sep 25, 2024

josephdviviano mentioned this pull request Sep 25, 2024

Add conditional LogZ calculation #150

Open

saleml reviewed Sep 25, 2024

View reviewed changes

josephdviviano added 8 commits October 1, 2024 12:10

API changes adapted

96b725c

added conditioning to all gflownets

5cd32a7

both trajectories and transitions can now store a conditioning tensor

877c4a0

input_dim setting is now private

279a313

added exception handling for all estimator calls potentially involvin…

65135c1

…g conditioning

API change -- n vs. n_trajectories

b4c418c

change test_box target value

738b062

API changes

4434e5f

josephdviviano marked this pull request as ready for review October 1, 2024 16:34

josephdviviano added 4 commits October 1, 2024 13:16

hacky fix for problematic test (added TODO)

851e03e

working examples for all 4 major losses

5152295

added conditioning indexing for correct broadcasting

1d64b55

added a ConditionalScalarEstimator which subclasses ConditionalDiscre…

348ee82

…tePolicyEstimator

josephdviviano added 3 commits October 4, 2024 16:09

added modified DB example

9120afe

conditioning added to modified db example

f59f4de

black

c5ef7ea

saleml reviewed Oct 5, 2024

View reviewed changes

josephdviviano added 6 commits October 8, 2024 23:29

reorganized keyword arguments and fixed some type errors (not all)

d67dfd5

reorganized keyword arguments and fixed some type errors (not all)

d56a798

added typing and a ConditionalScalarEstimator

db8844c

added typing

e03c03a

typing

6b47e06

typing

988faf0

josephdviviano mentioned this pull request Oct 9, 2024

conditioning example should do something #190

Open

added kwargs

f2bbce3

saleml reviewed Oct 17, 2024

View reviewed changes

saleml approved these changes Oct 17, 2024

View reviewed changes

josephdviviano added 2 commits October 24, 2024 15:20

renamed torso to trunk

eb13a2d

renamed torso to trunk

fd3d9dc

josephdviviano merged commit d2d959e into master Oct 24, 2024
3 checks passed

josephdviviano deleted the conditional_gfn branch October 24, 2024 19:22

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

conditional gfn #188

conditional gfn #188

josephdviviano commented Sep 25, 2024 •

edited

Loading

josephdviviano commented Sep 25, 2024

saleml Sep 25, 2024

josephdviviano Sep 25, 2024

josephdviviano Sep 25, 2024

saleml Sep 25, 2024

saleml Sep 25, 2024

josephdviviano Sep 25, 2024

saleml Sep 25, 2024

josephdviviano Oct 1, 2024

saleml commented Sep 25, 2024

saleml left a comment

saleml Oct 5, 2024

josephdviviano Oct 8, 2024 •

edited

Loading

saleml Oct 8, 2024

josephdviviano Oct 9, 2024 •

edited

Loading

saleml Oct 5, 2024

josephdviviano Oct 9, 2024

saleml Oct 5, 2024

josephdviviano Oct 9, 2024

josephdviviano Oct 9, 2024

saleml Oct 5, 2024

josephdviviano Oct 9, 2024

saleml Oct 5, 2024

josephdviviano Oct 9, 2024

saleml Oct 17, 2024

saleml Oct 17, 2024

josephdviviano Oct 24, 2024

saleml commented Oct 17, 2024

conditional gfn #188

conditional gfn #188

Conversation

josephdviviano commented Sep 25, 2024 • edited Loading

josephdviviano commented Sep 25, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

saleml commented Sep 25, 2024

saleml left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

josephdviviano Oct 8, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

josephdviviano Oct 9, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

saleml commented Oct 17, 2024

josephdviviano commented Sep 25, 2024 •

edited

Loading

josephdviviano Oct 8, 2024 •

edited

Loading

josephdviviano Oct 9, 2024 •

edited

Loading