Demo: Can an Unmasked Transformer Cheat?

A causal language model must predict the next token without seeing it. We train on independent random bits, where the past contains no information about the future, so any apparent advantage is easy to diagnose.

Two attention models share their initialization, batches, and optimizer. Their only difference is the causal mask.

import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F

plt.style.use("../figure_style.mplstyle")
torch.manual_seed(21)

length, width = 12, 24

The Attention Model

The module groups the trainable weights. Its forward method follows query, key, and value multiplication directly. Each batch contains several sequences, so transposing the last two axes transposes each key matrix.

class NextBitAttention(nn.Module):
    def __init__(self):
        super().__init__()
        self.token = nn.Embedding(2, width)
        self.position = nn.Parameter(torch.randn(1, length, width) / np.sqrt(width))
        self.Wq = nn.Linear(width, width, bias=False)
        self.Wk = nn.Linear(width, width, bias=False)
        self.Wv = nn.Linear(width, width, bias=False)
        self.head = nn.Linear(width, 2)

    def forward(self, bits, causal):
        n = bits.shape[1]
        h = self.token(bits) + self.position[:, :n]
        Q = self.Wq(h)
        K = self.Wk(h)
        V = self.Wv(h)
        scores = Q @ K.transpose(-2, -1) / np.sqrt(width)
        if causal:
            mask = torch.triu(torch.ones(n, n, dtype=torch.bool), diagonal=1)
            scores = scores.masked_fill(mask, float("-inf"))
        attention = F.softmax(scores, dim=-1)
        output = attention @ V
        return self.head(output), attention

Both copies start with identical weights and use Adam with the same learning rate.

initial = NextBitAttention().state_dict()
models = {"Causal mask": NextBitAttention(), "No mask": NextBitAttention()}
for model in models.values():
    model.load_state_dict(initial)
optimizers = {}
for name, model in models.items():
    optimizers[name] = torch.optim.Adam(model.parameters(), lr=3e-3)

Question: Can Low Training Loss Hide Information Leakage?

Random guessing has cross-entropy \(\log 2\). We compare teacher-forced loss, average attention weights, and accuracy when each model receives only a prefix. If the unmasked model learns to look ahead during training, decide which measurements will expose it.

steps, batch_size = 800, 128
history = {name: [] for name in models}
for step in range(steps + 1):
    batch = torch.randint(0, 2, (batch_size, length))
    for name, model in models.items():
        logits, _ = model(batch, causal=(name == "Causal mask"))
        next_bit_logits = logits[:, :-1].reshape(-1, 2)
        next_bits = batch[:, 1:].reshape(-1)
        loss = F.cross_entropy(next_bit_logits, next_bits)
        if step % 20 == 0:
            history[name].append(loss.item())
        if step == steps:
            continue
        optimizers[name].zero_grad()
        loss.backward()
        optimizers[name].step()

Prediction from a Prefix

At generation time the next bit is absent. We remove it from the input and score the prediction at the final visible position.

check = torch.randint(0, 2, (512, length))
with torch.no_grad():
    attention = {}
    prefix_accuracy = {}
    for name, model in models.items():
        causal = name == "Causal mask"
        _, weights = model(check, causal=causal)
        attention[name] = weights.mean(dim=0).numpy()
        prefix = check[:, :-1]
        prefix_logits, _ = model(prefix, causal=causal)
        prediction = prefix_logits[:, -1].argmax(dim=-1)
        prefix_accuracy[name] = (prediction == check[:, -1]).float().mean().item()

Compare the loss curves with where each model places its attention. Positions above the diagonal are future tokens.

fig, axes = plt.subplots(1, 3, figsize=(9, 3.4), layout="constrained")
x = np.arange(len(history["Causal mask"])) * 20
for name, linestyle in [("Causal mask", "-"), ("No mask", "--")]:
    axes[0].plot(x, history[name], linestyle=linestyle, linewidth=1.5, label=name)
axes[0].axhline(np.log(2), linestyle=":", linewidth=1,
                label="Random-bit limit")
axes[0].set_xlabel("Optimization Step")
axes[0].set_ylabel("Teacher-Forced Loss")
axes[0].legend(frameon=False)
for ax, name in zip(axes[1:], ["Causal mask", "No mask"]):
    im = ax.imshow(attention[name], vmin=0, vmax=1, cmap="Blues")
    ax.set_xlabel("Key Position")
    ax.set_title(name)
    ax.set_ylabel("Query Position")
fig.colorbar(im, ax=axes[1:], fraction=0.045, pad=0.03, label="Mean Attention Weight")
plt.show()

for name in models:
    print(f"{name}: final teacher-forced loss {history[name][-1]:.3f}, "
          f"prefix-only accuracy {prefix_accuracy[name]:.1%}")

Causal mask: final teacher-forced loss 0.693, prefix-only accuracy 53.3%
No mask: final teacher-forced loss 0.000, prefix-only accuracy 52.0%

Answer. The unmasked model lowers teacher-forced loss by attending above the diagonal, where the target bit is already visible. That advantage disappears when the future token is absent at generation time. The causal model stays near the random-bit limit because the task is genuinely unpredictable from the past.

The mask defines the information available to each prediction. Changing the mask changes which conditional prediction problem the model solves.