[Inactive] Changes for basic LLaDA style diffusion masking support #238
Conversation
PierreAndreNoel
left a comment
PierreAndreNoel
left a comment
There was a problem hiding this comment.
This is just quick feedback as I am very busy with other things, but please remind me to come back here next week and I'll dig deeper in.
|
|
||
| t = torch.rand(batch_size, device=device) | ||
|
|
||
| p_mask = (1 - diffusion_config.epsilon) * t + diffusion_config.epsilon |
There was a problem hiding this comment.
Some questions/thoughts (I am just browsing quickly, and I am not looking at the paper right now):
- Why is the lower bound
epsilonand the upper boundmax_mask_prob? - My guts tell me you never want the mask probability to be exactly 1, for the same kind of reasons you don't want it to be exactly 0.
- This approach using
torch.minwill put a discrete probability forp_maskto be exactlymax_mask_prob.
There was a problem hiding this comment.
Are you saying this coz we could have many timesteps with the exact masking level set to max_mask_prob? So are you suggesting some soft clipping instead of a hard upper bound?
|
|
||
| masked_indices = torch.rand((batch_size, seq_len), device=device) < p_mask | ||
|
|
||
| if diffusion_config.pad_prob > 0: |
There was a problem hiding this comment.
Meta: I currently can't comment about padding; it will have to wait for next week, as I need to re-read the paper better (our own work doesn't do padding).
There was a problem hiding this comment.
Is this to include variable length sequences for 1% of the data?
| p_mask = torch.min(p_mask, torch.tensor(diffusion_config.max_mask_prob)) | ||
| p_mask = p_mask[:, None].expand(-1, seq_len) | ||
|
|
||
| masked_indices = torch.rand((batch_size, seq_len), device=device) < p_mask |
| attention_mask = torch.ones((batch_size, 1, seq_len, seq_len), device=device, dtype=torch.bool) | ||
| else: | ||
| # Causal attention | ||
| attention_mask = torch.ones((batch_size, 1, seq_len, seq_len), device=device, dtype=torch.bool).tril_() |
There was a problem hiding this comment.
My understanding is that you never want such a triangular causal attention, as this would give a strictly worse model than an autoregressive model.
Suppose that, at inference, tokens are unmasked in the order (4, 2, 3, 0, 1). Token 4 is unmasked first, but this triangular matrix prevents all other tokens from ever "seeing" it.
What is the closest case that makes sense to me would be to permute the rows and columns of the triangular matrix using (4,2,3,0,1), so that token 2 can see token 4, token 3 can see tokens 2 and 4, etc.
There was a problem hiding this comment.
Yes, permuted rows and columns makes sense - so we can preserve the order in which it was unmasked. I will update this.
I guess this idea is similar to this paper? https://arxiv.org/abs/1906.08237
| kwargs['masked_indices'] = masked_indices | ||
| kwargs['p_mask'] = p_mask | ||
|
|
||
| if self._config.diffusion.bidirectional_attention: |
There was a problem hiding this comment.
You may want a string instead of a boolean, as there are many possible attention choices (e.g., blocks) that may come up. Also see the next comment below.
There was a problem hiding this comment.
I agree, will change this!
| masked_p = p_mask[masked_indices] | ||
|
|
||
| # Compute MLM loss | ||
| loss, grad = cross_entropy_forward_backward( |
There was a problem hiding this comment.
Is the loss already divided by masked_p? https://github.com/ML-GSAI/SMDM/blob/583aa4716d17728dbb825aec6c24a121164d616a/pretrain/train_mdm.py#L274
|
Is this still relevant? |
jlamypoirier
changed the title
✨ Description
Cleaned up the code a bit:
Of course still a WIP but feel free to leave comments and suggestions
These are changes to address this PR: #208 (comment)