Changes for basic LLaDA style diffusion masking support #238
Conversation
|
|
||
| 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.
|
|
||
| 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).
| 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.
| 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.
✨ 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)