Add precursor_cond support in diffusion model for precursor feature masks integration #31
Conversation
|
@singjc is above fine? |
There was a problem hiding this comment.
Great start so far. The code itself looks fine, except unet1d needs to be updated accordingly with the added precursor_cond mask.
With that being said though, we still need to think how exactly we want to use the precursor feature masks. Do we want to add them as additional conditioning signals, or do we want to replace the MS1 1D conditioning signal with the precursor feature masks for the attention condition?
Before merging, we would probably want to run tests first to see if there is any gain/benefit to the precursor feature masks and to make sure that everything still works as expected.
dquartic/model/model.py
Outdated
| @@ -241,7 +241,7 @@ def q_sample(self, x_0, t, noise=None): | |||
|
|
|||
| return sqrt_alpha_bar_t * x_0 + sqrt_one_minus_alpha_bar_t * noise | |||
|
|
|||
| def p_sample(self, x_t, t, init_cond=None, attn_cond=None): | |||
| def p_sample(self, x_t, t, init_cond=None, attn_cond=None, precursor_cond=None): | |||
There was a problem hiding this comment.
Would we want to have precursor_cond as a separate conditioning or would we replace the current MS1 that's passed to attn_cond with the precursor mask?
@LLYX what do you think?
dquartic/model/model.py
Outdated
| @@ -268,12 +269,12 @@ def p_sample(self, x_t, t, init_cond=None, attn_cond=None): | |||
|
|
|||
| if self.pred_type == "eps": | |||
| # Predict noise | |||
| eps_pred = self.model(x_t, t_tensor, init_cond, attn_cond) | |||
| eps_pred = self.model(x_t, t_tensor, init_cond, attn_cond, precursor_cond) | |||
There was a problem hiding this comment.
The current backbone model is the unet1d model, which the forward method currently accepts forward(self, x, time, init_cond=None, attn_cond=None).
dquartic/model/model.py
Outdated
| # Compute x_0 prediction | ||
| x0_pred = (x_t - sqrt_one_minus_alpha_bar_t * eps_pred) / sqrt_alpha_bar_t | ||
| elif self.pred_type == "x0": | ||
| # Predict x_0 directly | ||
| x0_pred = self.model(x_t, t_tensor, init_cond, attn_cond) | ||
| x0_pred = self.model(x_t, t_tensor, init_cond, attn_cond, precursor_cond) |
dquartic/model/model.py
Outdated
| @@ -356,7 +361,7 @@ def train_step(self, x_0, ms2_cond=None, ms1_cond=None, noise=None, ms1_loss_wei | |||
|
|
|||
| if self.pred_type == "eps": | |||
| # Predict noise | |||
| eps_pred = self.model(x_t, t, ms2_cond, ms1_cond) | |||
| eps_pred = self.model(x_t, t, ms2_cond, ms1_cond, precursor_cond) | |||
There was a problem hiding this comment.
Same comment as previous above, relating to the current unet1d models expected inputs
dquartic/model/model.py
Outdated
| @@ -371,7 +376,7 @@ def train_step(self, x_0, ms2_cond=None, ms1_cond=None, noise=None, ms1_loss_wei | |||
| ) | |||
| elif self.pred_type == "x0": | |||
| # Predict x0 | |||
| x0_pred = self.model(x_t, t, ms2_cond, ms1_cond) | |||
| x0_pred = self.model(x_t, t, ms2_cond, ms1_cond, precursor_cond) | |||
There was a problem hiding this comment.
Same comment as previous above, relating to the current unet1d models expected inputs
|
Hi @singjc, I've implemented the feature mask integration using the replacement approach we discussed. Changes made:
This implementation follows the simpler approach of replacing MS1 with feature masks rather than adding them as an additional conditioning signal. The changes maintain compatibility with the existing UNet1d architecture by ensuring the feature masks will have the same tensor shapes and normalization as the MS1 data they replace. Please let me know if you'd like any modifications to this implementation approach. |
Hi @singjc,
This pull request updates the dquartic/model/model.py file to integrate precursor feature masks as an additional conditioning signal in the diffusion model. The following changes were made:
def p_sample:
Added a new precursor_cond parameter.
Passed precursor_cond to the model’s forward invocation.
Updated the docstring to include details about the new parameter.
def sample:
Added precursor_cond as an input argument.
Normalized precursor_cond in the same way as other conditioning signals.
Forwarded precursor_cond to the p_sample method.
Updated the docstring accordingly.
def train_step:
Added a precursor_cond parameter.
Normalized precursor_cond and included it in the model's forward invocation.
Updated the docstring accordingly.
These changes allow the model to receive precursor feature masks as additional conditioning during both training and sampling. This integration is part of the ongoing efforts for the GSoC issue #20 "Generate MS1 and MS2 Peptide Feature Masks for Targeted Deconvolution."
Please review the changes and let me know if further modifications are required.