Technical Summary: Gumbel Machine for Counterfactual Student Writing Generation

Problem Statement

Effective educational feedback often relies on providing examples of high-quality work. However, standard model examples frequently differ significantly from a student's current output in style, approach, or personal tendencies, making it difficult for learners to emulate the improvements. The ideal feedback mechanism is a counterfactual version of the student's work: an improved version that remains similar to the original while addressing specific deficiencies.

Existing automated approaches using Large Language Models (LLMs) face a central design challenge: balancing task performance (generating a text that meets a specific rubric score) with similarity (maintaining fidelity to the original student text). Current methods often rely on:

1. Joint Optimization: Encoding the trade-off in model weights, which limits flexibility.
2. Decomposition: Identifying and replacing specific words, which implicitly enforces similarity but often fails to capture nuanced, ordinal improvements required in education.
3. Task-Specific Training: Many approaches require retraining or system redesign for new tasks, limiting generalizability.
4. Nominal vs. Ordinal: Most existing counterfactual generation targets nominal labels (e.g., sentiment, topic), whereas educational scoring is typically ordinal (e.g., rubric scores 1–5), requiring nuanced adjustments rather than binary category shifts.

Methodology: The Gumbel Machine

The authors propose the Gumbel Machine, a modular approach that decouples similarity control from task alignment. The core innovation is $\beta$-Hindsight control, a model-agnostic, inference-time decoding mechanism.

1. Theoretical Foundation: Gumbel-Max Trick

The method leverages the Gumbel-Max trick, which allows sampling from a categorical distribution by adding Gumbel noise to logits and selecting the maximum. Crucially, the noise values ($g_v$) can be stored to preserve the stochastic state of the sampling process.
$$y = \arg \max_{v \in V} (\ell_v + g_v)$$
where $\ell_v$ are logits and $g_v \sim \text{Gumbel}(0, 1)$.

2. $\beta$-Hindsight Control Algorithm

The process operates in two stages:

• Stage 1: Noise Recovery (Hindsight Sampling): Given a reference factual text $y$ generated by an LLM, the system approximates the random process to recover the Gumbel noise values ($g_t$) that would have led to the observed tokens. This involves sampling from truncated Gumbel distributions to satisfy the condition that the observed token was the maximum.
• Stage 2: Replay with Intervention: To generate a counterfactual $y'$, the system modifies the input prompt (e.g., changing the target rubric score) and replays the recovered noise values. The noise is scaled by a hyperparameter $\beta$:
  $$y'_t = \arg \max_{v \in V} (\ell'_v + \beta \cdot g_t[v])$$
  Here, $\beta$ acts as a tunable control signal: higher values induce outputs more similar to the reference, while lower values allow the model to prioritize the new target score.

3. Task Validity Alignment

To ensure the generated counterfactuals are valid (i.e., achieve the target rubric score), the authors combine $\beta$-Hindsight with a preference optimization strategy:

• Supervised Fine-Tuning (SFT): Adapts the model to the specific student writing style and task format.
• Direct Preference Optimization (DPO): Aligns the model with the counterfactual task by training on preference pairs where a text achieving the target score is preferred over one with a different score.

Key Contributions

1. Decoupled Control: The Gumbel Machine separates similarity control (via $\beta$-Hindsight) from task validity (via instruction alignment and DPO), allowing for flexible adjustment of the similarity-validity trade-off without retraining.
2. Tunable Similarity via Recovered Noise: The paper demonstrates that recovered Gumbel noise serves as an effective, tunable control signal for similarity, outperforming token-level heuristics.
3. Ordinal Task Adaptation: The approach is specifically designed for ordinal scoring tasks in education, moving beyond the nominal label focus of prior counterfactual generation work.
4. Modular Architecture: The system is compatible with both open-weight and proprietary LLMs and does not require task-specific classifiers or control codes.

Experimental Results

The authors evaluated the approach on two datasets: CLASSE (student summaries) and DREsS (EFL essays), using both open-weight (Llama-3, Qwen) and proprietary (Claude) models.

• Performance vs. Baselines: The Gumbel Machine (GM) outperformed baselines including Minimal-Edit (ME), In-Context Learning (ICL), and Identify-and-Replace (I&R).
  • On CLASSE, GM achieved the highest validity (0.668) while maintaining strong similarity (0.391), whereas I&R had high similarity but low validity.
  • On DREsS, GM outperformed all baselines on both metrics.
• Cross-Model Generalization: GM translated effectively across different model families (Llama, Qwen, Claude), often matching or exceeding the performance of larger proprietary models.
• Human Evaluation: Teachers evaluated GM rewrites against ME baselines.
  • GM rewrites were rated as significantly more similar to the original student work (62% win rate in contested cases, $p=0.033$).
  • When GM rewrites were judged more similar, teachers also preferred them for pedagogical utility (75% win rate).
• Ablation Studies:
  • $\beta$-Hindsight alone improved similarity but failed to ensure validity.
  • SFT alone drastically improved similarity but collapsed validity.
  • DPO was necessary to realize counterfactuals that balanced both axes. All three components (Recovery, SFT, DPO) were required for optimal results.

Significance and Claims

The paper claims that the Gumbel Machine offers a practical, flexible solution for generating counterfactual student work. Its primary significance lies in:

• Pedagogical Relevance: By generating feedback that is both rubric-consistent and stylistically similar to the student's original work, the system addresses the "mismatch" problem of traditional model examples, potentially making feedback more actionable for learners.
• Technical Innovation: It introduces a novel use of recovered stochastic states (Gumbel noise) as a control mechanism, demonstrating that preserving the randomness of the sampling process is more effective for similarity control than simple token-level constraints.
• Scalability: The modular nature of the approach suggests it can be adapted to various educational domains without the heavy retraining costs associated with previous task-specific systems.

The authors remain modest regarding limitations, noting that the approach assumes an autoregressive approximation of student writing (which involves planning and revision) and relies on automated scoring models that may introduce biases. They emphasize that the system is intended as an illustrative demonstration rather than an authoritative correction to preserve student agency.