about decoding topk_masking

[violet-sto]

Hi

Thanks for your excellent work. I have a question about the rate schedule for topk_masking.

As described in the appendix, "To ensure that the degree of noise decreases as the generation process proceeds, we schedule k to increase from 1 to N monotonically as the diffusion step t goes from T to 1." However, in the code (

reparam-discrete-diffusion/discrete_diffusion/discrete_diffusions/discrete_diffusion_base.py

Line 177 in 26ee286

lowest_k_mask = topk_masking(_scores_for_topk, cutoff_len, stochastic=False)

), the masked k tokens with the lowest confidence instead of the highest. Are there any inconsistencies here?

Best regards

[LZhengisme]

Hi there,

Thanks for reaching out! In

reparam-discrete-diffusion/discrete_diffusion/discrete_diffusions/utils.py

Lines 20 to 24 in 26ee286

	sorted_index = _scores.sort(-1)[0]
	cutoff = sorted_index.gather(dim=-1, index=cutoff_len) + 1e-10
	# cutoff_len = k -> select k + 1 tokens
	masking = _scores < cutoff
	return masking

, the topk_masking function actually returns a mask to indicate the unselected elements. This is essentially the inverse of selecting the highest elements; we implement this way to simplify subsequent calculations for denoising tokens.

Hope this clears things up!