Doge: Stopping LLM Knowledge Theft With One Fine Tuned Layer

Analysis by the aitrendblend editorial team · Pillar 2, Knowledge Distillation and Model Compression · 13 minute read

A well funded lab spends months and a serious pile of compute training a frontier language model. A smaller competitor never touches that budget. Instead it sends a few thousand prompts to the public API, saves the answers, and fine tunes a much cheaper model to mimic them. Within a week the copy performs close enough to the original that nobody wants to pay for the real thing. Researchers at UNC Chapel Hill and Arizona State University just published a way to make that copy come out broken, not by hiding the model or watching for theft after the fact, but by quietly retraining a single layer to sabotage anyone who tries.

The uncomfortable economics of a public API

Every serious LLM lab faces the same quiet threat. The moment a model answers questions through an API, its outputs become a training set for anyone willing to collect them. This is knowledge distillation, a technique originally proposed for compressing large models into smaller ones (Hinton et al., 2015), now repurposed as a cheap way to imitate a competitor’s product. Feed a smaller model enough input output pairs from a proprietary teacher and it can absorb a striking amount of that teacher’s capability, all without ever seeing a single weight or gradient of the original (Tramer et al., 2016). It is, as the authors put it, uncomfortably similar to learning a craft simply by watching someone else perform it well.

The tools labs currently have to fight this are mostly reactive. Watermarking schemes embed identifiable patterns in generated text so a lab can later prove theft occurred (Kirchenbauer et al., 2023). Fingerprinting techniques try to uniquely identify a model’s outputs after the fact. Both are useful for a lawsuit. Neither stops the copying while it is happening. Other active defenses exist, but most assume the attacker is mimicking the teacher’s internal probability distribution over the full vocabulary, the raw logits, which simply is not available to someone who can only call a public API and read back text.

Turning the problem into a game between two optimizers

What makes this paper interesting methodologically is how cleanly it frames anti distillation as an actual optimization problem rather than a vague design goal. Standard sequence level knowledge distillation trains a student S by minimizing a loss between its predictions and a teacher generated dataset of input output pairs. The defender’s job is the mirror image of that, framed as a nested optimization where an inner step models the student learning as well as it possibly can, and an outer step tunes the teacher to make that best case student as weak as possible.

Solving that nested optimization exactly is intractable, since it would require simulating a full student training run inside every step of teacher training. DOGe is the practical approximation the authors build instead, and the shortcut they take is clever precisely because it avoids ever training a real student during defensive training at all.

Why the last layer is enough

Rather than retraining the whole teacher model, DOGe touches only the LM head, the final linear layer that converts internal hidden states into a probability distribution over the vocabulary right before sampling. Everything else in the model stays frozen. That decision is not just about saving compute, although it does. It means the defense can be swapped in and out at serving time by loading a different small set of weights, and it means the enormous bulk of the model’s learned knowledge is never disturbed, only the very last step where that knowledge gets converted into a specific token probability.

An adversarial loss aimed at a stand in student

Training the teacher to fool an actual, unknown future attacker is impossible since that attacker’s model architecture is not known in advance. DOGe sidesteps this by training against one or two proxy student models the defender chooses ahead of time, models that share the teacher’s tokenizer so their logits are directly comparable. The total training loss combines two terms.

Maximizing a forward KL divergence can be numerically unstable in principle, since the loss can blow up if a student assigns near zero probability somewhere the teacher assigns real probability. The authors note three things that keep this from becoming a practical problem, that real softmax outputs almost never hit exact zero, that the supervised fine tuning term anchors the distribution to real answers throughout training, and that the trade off coefficient lambda itself needs to stay small enough for stability, which their own ablation studies confirm.

The trick that actually protects the answer

Pushing every single token away from what a student would predict sounds like it should also wreck the teacher’s own usefulness, and that is exactly the failure mode the paper’s most clever design choice avoids. DOGe splits every generated sequence into intermediate reasoning tokens and final answer tokens, then applies the adversarial pressure only to the former.

For models like DeepSeek R1 that already mark reasoning with explicit tags, this split is easy. For other models the authors fall back on regular expressions that look for answer formatting cues such as the word Answer followed by a colon like structure. Either way, the effect is the same. A user asking a question still gets the right number in the end. The path the model took to get there becomes noisier, more repetitive, occasionally strange, and much harder for a smaller model to learn a generalizable pattern from.

What the numbers actually showed

The headline experiments used two teacher models, DeepSeek-R1-7B and Qwen3-8B, each defended using GSM8K math problems as the training data and a pair of same family, smaller models as proxy students. They then measured accuracy on a held in dataset, GSM8K itself, and three held out datasets, MATH, ARC Challenge, and CommonsenseQA, comparing four groups, the original teacher, the defensive teacher, a student distilled from the original teacher, and a misled student distilled from the defensive teacher.

Two things jump out from that table. The defensive teacher column is mostly positive or barely negative, meaning DOGe is not degrading the product real users interact with, and in several cases it is quietly improving it, which the authors attribute to the adversarial training pushing the model toward more robust reasoning patterns even as it makes that reasoning harder for an outsider to copy. The misled student column tells the opposite story entirely, with drops as steep as 39 percentage points on CommonsenseQA when a Gemma-3-1b-it student tried to learn from the DeepSeek-R1-7B defensive teacher, a collapse the authors describe as roughly five times worse than the same student copying an undefended teacher.

The teacher got slightly better at answering questions. Anyone copying it got dramatically worse. Editorial synthesis of the paper’s central result

Turning the dials, what the ablations reveal

A defense is only useful in practice if its knobs behave predictably, and the paper runs a solid set of ablations to check exactly that. Sweeping the adversarial weight lambda across 1 times 10 to the negative 5, 3 times 10 to the negative 5, and 1 times 10 to the negative 4 traces out a clear Pareto frontier. At the smallest value the defensive teacher stays almost identical to the original but only mildly slows down a student. At the paper’s chosen default of 3 times 10 to the negative 5, teacher performance holds steady while student performance falls sharply, which the authors call the sweet spot. Push lambda to 1 times 10 to the negative 4 and both teacher and student collapse toward zero, since the adversarial pressure has grown strong enough to fight the supervised loss into instability rather than balance against it.

Two other findings matter for anyone thinking about actually deploying this. Adding a second proxy student model on top of the first changed defense effectiveness by less than one percentage point across every benchmark, while roughly doubling training cost, which tells you a single well chosen proxy captures most of what a whole ensemble would. And swapping the defensive training data from task specific GSM8K math problems to the broader, general purpose Tulu instruction dataset produced stronger student degradation across the board, at a small cost to how much the teacher improved on its own in domain math tasks, a genuine trade off a model owner gets to choose rather than one the method forces on them.

The paper also tested what happens against a student roughly the same size as the teacher itself, pitting an 8 billion parameter Llama-3.1-8B student against the 8 billion parameter Qwen3-8B teacher. The bigger student started from a stronger baseline after distillation, as you would expect, but it also suffered a larger relative collapse, dropping 20 to 50 percentage points across benchmarks compared with 8 to 18 points for the 1 billion parameter Llama-3.2-1B student. Defense effectiveness scaled up with student capacity rather than fading against a more capable copier, which is the opposite of what you might naively expect from a subtle, single layer intervention.

Looking inside the training dynamics

One of the more genuinely illuminating pieces of this paper is a direct visualization of the loss landscape under different values of lambda. With no adversarial term at all, the landscape is a smooth, well behaved basin with an obvious path to a clear minimum, the kind of picture you would expect from ordinary supervised fine tuning. Add a small adversarial weight and the basin develops mild ripples but keeps a dominant path toward the minimum, meaning the model is still perfectly trainable, just nudged. Push the coefficient up toward 1 times 10 to the negative 3 and the landscape turns genuinely rough, with sharp gradients and what look like competing local minima, which lines up exactly with the diverging, unstable training curves the authors observed once lambda reached 1 times 10 to the negative 4 in their actual runs. The visualization is not just decoration. It is a real explanation for why the earlier Pareto frontier ablation bends the way it does.

What the poisoned reasoning actually looks like

The case study section is worth dwelling on because it shows the mechanism in plain text rather than in a chart. In one successful example, a defensive teacher answering a question about download time keeps its final answer correct while filling its reasoning with repeated phrases like checking again to make sure no mistake was made, triple checking, and reaffirming the same conclusion several times over, the kind of padded, low confidence hedging a student model would struggle to distill anything generalizable from. In a rarer failure case, the reasoning trace degrades further still, briefly mixing in fragments of a different language before recovering to the correct boxed answer, evidence that the adversarial pressure occasionally pushes the reasoning process somewhere genuinely strange rather than just verbose.

To check whether this strangeness would tip off a careful observer, the authors ran a separate evaluation using Gemini-1.5-Pro as an automated judge, asking it to label a sample of reasoning traces as natural or unnatural sounding. Across four benchmarks, the share of outputs the judge flagged as unnatural ranged from 19 percent on ARC up to 31 percent on MATH, meaning a clear majority of the defensively generated reasoning passed as ordinary even under direct scrutiny.

Where this defense could still be beaten, or misused

The bigger idea worth taking away

Strip away the specific benchmarks and DOGe is really proposing a shift in where model protection lives. Watermarking and fingerprinting live downstream of the problem, in a forensic sense, telling you after the fact that theft happened. Decoding time interventions live at inference, adding cost and complexity every single time the model generates text. DOGe embeds the defense directly into the model’s weights, specifically into the smallest, cheapest to swap part of those weights, so protection becomes a property of the model itself rather than a wrapper around it. That is a meaningfully different design point, and the fact that it can be toggled by loading a different small file rather than retraining a whole model is exactly the kind of practical detail that decides whether a defense actually gets deployed or stays a paper result.

The conceptual move worth remembering is treating imitation resistance as an explicit dual objective optimization rather than a side effect of some other technique. Framing the problem that way, maximize your own usefulness while minimizing the usefulness of anything trained to copy you, is general enough that it plausibly extends beyond text models, to any service that exposes its outputs through an API a competitor could scrape and learn from. The specific mechanics here, an LM head and vocabulary logits, are text specific, but the underlying strategic framing is not.

What remains genuinely open is how this holds up against an adaptive attacker, someone who suspects a defense is active and adjusts their distillation strategy accordingly, perhaps by filtering out suspiciously repetitive reasoning before training a student on it. The paper does not test that scenario, and it is exactly the kind of arms race dynamic that tends to follow any new defense once it becomes widely known. Cross tokenizer robustness is another open question the authors flag rather than solve.

Read as a whole, this is a genuinely practical piece of applied security research for an economic problem the field has mostly handled with legal threats and after the fact detection. It will not end the broader tension between open competition and protecting real investment in frontier model training, and the authors are upfront that it should not be deployed carelessly against models meant for open research. What it does offer is a concrete, cheap, and measurable way for a model owner to make imitation expensive again, which is a rarer contribution in this space than another benchmark leaderboard entry.

Reproducing the core defense

The full experimental pipeline relies on large pretrained checkpoints and DeepSpeed for distributed training, but the actual defensive mechanism, the combined loss, the reasoning aware mask, and the training loop, is compact enough to implement directly, closely following Algorithm 1 in the paper’s appendix.

# doge_lm_head_defense.py
# Reproduces the DOGe defensive training loop from Li, Tan, Zhang, Qu, Liu
# and Chen, arXiv 2505.19504. Fine tunes only the teacher's LM head using
# a combined supervised and adversarial loss with reasoning aware masking.

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


class DefensiveLMHead(nn.Module):
    """Trainable final linear layer, Eq. 4.1 and 4.4. The base transformer
    that produces hidden states is assumed frozen and passed in externally."""

    def __init__(self, hidden_size: int, vocab_size: int):
        super().__init__()
        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.lm_head(hidden_states)  # (batch, seq_len, vocab_size) logits


def sft_loss(teacher_logits: torch.Tensor, target_ids: torch.Tensor) -> torch.Tensor:
    """Standard token level cross entropy against ground truth labels."""
    vocab_size = teacher_logits.shape[-1]
    return F.cross_entropy(
        teacher_logits.view(-1, vocab_size),
        target_ids.view(-1),
        reduction="none",
    ).view_as(target_ids)


def adversarial_loss(teacher_logits: torch.Tensor, proxy_logits_list: list, alpha: float = 2.0) -> torch.Tensor:
    """Negative average KL divergence between the teacher and each proxy
    student, temperature scaled, Eq. 4.2. Returns per token loss values,
    higher magnitude when the teacher is closer to the proxies."""
    teacher_log_probs = F.log_softmax(teacher_logits / alpha, dim=-1)
    teacher_probs = teacher_log_probs.exp()

    kl_terms = []
    for proxy_logits in proxy_logits_list:
        proxy_log_probs = F.log_softmax(proxy_logits / alpha, dim=-1)
        # KL(teacher || proxy), summed over vocab, per token
        kl = (teacher_probs * (teacher_log_probs - proxy_log_probs)).sum(dim=-1)
        kl_terms.append(kl)

    avg_kl = torch.stack(kl_terms, dim=0).mean(dim=0)
    return -avg_kl  # minimizing this maximizes the actual divergence


def reasoning_mask(is_reasoning_token: torch.Tensor) -> torch.Tensor:
    """Binary mask from Eq. 4.3. Expects a boolean tensor already computed
    from tag detection or regular expression based answer span finding,
    matching shape (batch, seq_len)."""
    return is_reasoning_token.float()


def doge_training_step(head: DefensiveLMHead, optimizer: torch.optim.Optimizer,
                       hidden_states: torch.Tensor, target_ids: torch.Tensor,
                       proxy_logits_list: list, is_reasoning_token: torch.Tensor,
                       lam: float = 3e-5, alpha: float = 2.0) -> dict:
    """One DOGe update, matching Algorithm 1. Combines the masked
    adversarial gradient with the full supervised gradient."""
    optimizer.zero_grad()

    teacher_logits = head(hidden_states)
    per_token_sft = sft_loss(teacher_logits, target_ids)
    per_token_adv = adversarial_loss(teacher_logits, proxy_logits_list, alpha=alpha)
    mask = reasoning_mask(is_reasoning_token)

    # Eq. 4.4, the adversarial term is masked, the SFT term is not
    per_token_total = per_token_sft + lam * mask * per_token_adv
    loss = per_token_total.mean()

    loss.backward()
    optimizer.step()

    return {
        "total_loss": loss.item(),
        "sft_loss": per_token_sft.mean().item(),
        "adv_loss": per_token_adv.mean().item(),
    }


def evaluate_divergence_from_proxies(head: DefensiveLMHead, hidden_states: torch.Tensor,
                                       proxy_logits_list: list, alpha: float = 2.0) -> float:
    """Reports the average KL divergence achieved, a proxy for how hard the
    current teacher would be to imitate through output based distillation."""
    with torch.no_grad():
        teacher_logits = head(hidden_states)
        divergence = -adversarial_loss(teacher_logits, proxy_logits_list, alpha=alpha)
    return divergence.mean().item()


def run_smoke_test():
    """Dummy hidden states, targets, and two proxy student logit sets,
    confirming the training step reduces SFT loss and increases divergence
    from the proxies over a few steps, without touching a real base model."""
    torch.manual_seed(7)
    batch, seq_len, hidden_size, vocab_size = 4, 16, 32, 100

    head = DefensiveLMHead(hidden_size, vocab_size)
    optimizer = torch.optim.AdamW(head.parameters(), lr=5e-5)

    hidden_states = torch.randn(batch, seq_len, hidden_size)
    target_ids = torch.randint(0, vocab_size, (batch, seq_len))
    proxy_logits_list = [torch.randn(batch, seq_len, vocab_size) for _ in range(2)]

    # mark the first half of each sequence as reasoning tokens, matching
    # the paper's split between thinking tokens and the final answer
    is_reasoning_token = torch.zeros(batch, seq_len, dtype=torch.bool)
    is_reasoning_token[:, : seq_len // 2] = True

    divergence_before = evaluate_divergence_from_proxies(head, hidden_states, proxy_logits_list)

    for step in range(50):
        stats = doge_training_step(
            head, optimizer, hidden_states, target_ids,
            proxy_logits_list, is_reasoning_token, lam=3e-5, alpha=2.0,
        )
        if step % 10 == 0:
            print(f"Step {step}, total {stats['total_loss']:.4f}, sft {stats['sft_loss']:.4f}, adv {stats['adv_loss']:.4f}")

    divergence_after = evaluate_divergence_from_proxies(head, hidden_states, proxy_logits_list)

    print(f"Divergence from proxies before {divergence_before:.4f}, after {divergence_after:.4f}")
    assert divergence_after > divergence_before
    print("Smoke test passed.")


if __name__ == "__main__":
    run_smoke_test()

Frequently asked questions

Explore more from this pillar