Implement a small end-to-end masked diffusion LM sampler with confidence-based remasking a

Super-hard

Diffusion & Non-Autoregressive Language Models4 / 4

Chapter 224 / 4

Diffusion & Non-Autoregressive Language Models

Implement a small end-to-end masked diffusion LM sampler with confidence-based remasking a

Super-hard

Implement a small end-to-end masked diffusion LM sampler with confidence-based remasking and compare its output/latency to a greedy AR baseline on a toy vocabulary.

Target interface

target(L)

Complete the starter code in the editor. Any correct approach is accepted.

Hints0 / 2

Reference solution

Your own code stays in the editor.

solution.pyPython

Local draft

import time
import torch
import torch.nn.functional as F

def target(L):
    raise NotImplementedError

class Oracle:

    def __call__(self, x):
        raise NotImplementedError

@torch.no_grad()
def diffusion_sample(model, L, steps, mask_id=MASK):
    raise NotImplementedError

@torch.no_grad()
def ar_sample(model, L, mask_id=MASK):
    raise NotImplementedError

Local draft

⌘/Ctrl↵to submit

AI review

Ready when you are

Submit your solution for a score and feedback on correctness, edge cases, and clarity. Any correct approach is accepted.

Chapter 224 / 4

Diffusion & Non-Autoregressive Language Models

Implement a small end-to-end masked diffusion LM sampler with confidence-based remasking a

Super-hard

Implement a small end-to-end masked diffusion LM sampler with confidence-based remasking and compare its output/latency to a greedy AR baseline on a toy vocabulary.

Target interface

target(L)

Complete the starter code in the editor. Any correct approach is accepted.

Hints0 / 2

Reference solution

Your own code stays in the editor.