Gumbel Counterfactual Generation From Language Models

Thank you for your constructive feedback and for recognizing the importance of our work in advancing the true counterfactual generation from LMs. We appreciate the strengths highlighted, especially regarding our theoretical contributions and the insightful experimental findings. We would like to address the concerns and provide clarifications to improve the understanding of our work.

Comparison with existing methods

The reviewer suggests including a comparison with existing methods for counterfactual generation. We would like to clarify that, to the best of our knowledge, there is no existing method that proposes a framework for generating causally correct counterfactuals in the context of language models. Most prior approaches, such as linear steering or knowledge editing, intervene on model behavior without addressing the core challenge of identifying causal counterfactuals by conditioning on the same underlying noise that produced the original sentence. In the existing literature people vaguely refer to these as “counterfactuals”, while, as we point out, they actually correspond to interventions. We use several such interventions in our experimental evaluation and generate true counterfactual strings from them. The main contribution of our work is the introduction of a method that generates true counterfactuals by explicitly modeling and controlling for the sampling noise, an ability that doesn’t currently exist. This is why we focused on evaluating our proposed method.

Question 1: Applicability of the method to beam search decoding

The reviewer asks whether our proposed method can work with beam search decoding, which is commonly used in LLMs. Please note that beam search is a deterministic decoding strategy that attempts to locate the sequence with the highest cumulative probability (i.e., the approximate argmax of the model's output distribution). It is not a sampling method. Unlike stochastic sampling methods, beam search does not introduce any exogenous noise into the generation process. Our method specifically targets scenarios involving stochastic sampling, where the generation process involves a random component (e.g., multinomial sampling) driven by exogenous noise variables. Our method can be easily extended to other stochastic sampling techniques, such as nucleus sampling or top-k sampling, as we mention in a footnote in the paper. These methods can be incorporated by applying a deterministic function to the logits (e.g., truncating the probability distribution) before the sampling step. The crucial point is that, as long as there is some form of stochastic sampling involved, we can model the exogenous noise and apply our counterfactual generation algorithm. In the case of beam search, there is no exogenous noise to model, and thus, our algorithm for calculating counterfactuals is not applicable. We will clarify this distinction in the revised version of the paper.

Question 2: Effect of temperature parameter in the gumbel-max

The reviewer asks how the temperature parameter in the Gumbel-max trick affects the quality of the generated counterfactual sentences. The temperature parameter is a deterministic part of the forward pass that scales the logits before sampling. This scaling influences the model's output distribution, affecting the text quality and diversity by controlling the randomness of the token selection. However, this scaling is part of the deterministic transformation applied to the logits, and as such, it does not affect the exogenous noise component that our method aims to model and disentangle. Our counterfactual generation method focuses on controlling the exogenous stochastic noise in the sampling process. Since the temperature only modifies the logits deterministically, our method is invariant to the choice of temperature. While different temperatures can influence the overall quality and diversity of the generated text, this is orthogonal to our main objective of deriving a counterfactual given a deterministic model. The exogenous noise remains the same regardless of the temperature setting.

Thanks for the in-depth review. We wanted to answer the three questions before we respond with more because we think we can address all your concerns in detail, but we want to be sure we understood it all.

1. Our method is invariant to the exact method used to sample from the softmax in a specific sense. Non-identifiability when sampling from Categorical arises, for example, from permutation non-invariance. See, for example, Oberst and Sontag (2019) Secton 3.1 [https://proceedings.mlr.press/v97/oberst19a/oberst19a.pdf]. We can give formal proof of this, which appears to be missing from the above reference. Thus, we chose the Gumbel-max trick for convenience as it is a simple, permutation-invariant sampling scheme. Would our spelling out this equivalence assuage some of your concerns? The reason we rule out permutation-dependent schemes, e.g., what we would call roulette sampling, is that the permutation is arbitrary.

To be explicit, the result we will add, which we think will dramatically improve the paper is this. If one wishes to assume the sampling process from the Categorical is permutation-invariant, i.e., the order we assign elements of the categorical an integer does not matter, and that sampling is useing additive noise, then the model is identifiable.

2. This is a choice to reflect standard LM notation. We can change it it was confusing.

3. All of our experiments use hindsight sampling. Roughly speaking, the experimental design works like this: a) Input a string b) Apply an intervention on the network c) Sample a counterfactual string using Algorithm 1.

So, our method is the first that can show how individual strings change under interventions, e.g., MEMIT.

We specifically appreciated many of your comments about GSEMs. One quick note -- our GSEM is acyclic and well-founded, so they have one solution. (An unrolled neural network-based language model is acyclic and well-founded.) We only apply them because we have an infinite number of variables and SEMs do not apply. Apologies that we did not make this point clear. Also, we are not tied to GSEMs, as the reviewer infers, would it help the presentation to switch to causal dynamic Bayesian networks? We view our contribution as the first algorithm to give string-level counterfactuals under the model. We can elaborate more on why we call these "true" in the title.

Following our initial comment, we provide here a detailed response. We thank you for your comments, which helped as to improve the presentations of several keys elements in the revised manuscript (marked in red).

Experiments and Algorithms: We want to emphasize that in all experiments, we used exactly the algorithm for hindsight Gumbel generation, ensuring that the experiments align perfectly with the theoretical framework presented in the paper. If any part of the presentation gave rise to a different interpretation, please let us know so we can address it.

Why use lengths of prefixes for the evaluation: please note that we have also reported the cosine similarity under a strong, widely used encoder, which gives a notion of “semantic” divergence”. We agree, however, that we can report additional metrics as well. We repeat the experiment with the edit distance. The results are available here: https://imgur.com/a/8skFOe3 The trends are similar to the currently reported results. We will add them to the final version of this work.

Identifiability: We thank you for highlighting the question of identifiability. In the revised PDF, we have added a new section in the appendix that provides an in-depth discussion of the identifiability of the counterfactual distribution and justifies our construction. For your convenience, these additions are marked in red. In short, we observe that an alternative approach, inverse CDF sampling, is sensitive to arbitrary decisions, such as the mapping of events to integers, making it less natural. Previous work has highlighted the counterfactual stability of the Gumbel mechanism [2], although it is not the only counterfactually stable mechanism. Importantly, if we assume a sampling process based on taking the argmax after an additive noise (a “Thurstone model” in decision theory[1]), the Gumbel distribution uniquely leads to the softmax function, which is a key component of the causal graph. This ensures the resulting counterfactual distribution is identifiable (again, under this modeling choice). In the updated manuscript, we refer to the proof for this claims in Appendix B, based on results in [3].

Proposition 2.1: following your suggestion, we now present the result of identifying LMs as GSEMs as a construction, and we explicitly note that the causal model in our case is acyclic.

Notation: Thank you for highlighting the potential confusion. We have revised the notation in several places throughout the paper. Additionally, we explain the term “truncated probability distribution” to avoid ambiguity. This term was originally adopted from [4], a highly cited work. We want to clarify that our reference to the blog refers specifically to the implementation, which is significantly more efficient than naive rejection sampling and is, in fact, equivalent to conditional inverse CDF sampling.

Q1: please refer to toe discussion in identifiability, here and in the paper.

Q2: Regarding E vs. h, we found it is relevant to present them separately because this highlights how the latent encodings give rise to the logits (from which we sample). We agree this differentiation is not essential, and can change it if you think it introduces confusion.

Q3: MEMIT and similar techniques are methods for inducing interventions in the base model. For example, MEMIT applies a low-rank update to specific weight matrices in the FFN layers. After performing this intervention, we use Algorithm 1 to generate string counterfactuals, aiming to answer the following question: How would a specific sentence (originally generated by the base model) appear if it had been produced by the model after the MEMIT intervention? This is done by generating from the new model, under the noise derived from the base model. In all our experiments, we apply exactly the theory we propose—particularly Algorithm 1—to derive string counterfactuals under various intervention techniques. The process is straightforward: (1) generate a sentence from the base model; (2) hindsight-sample Gumbel noise conditioned on this generated sentence; and (3) use the same Gumbel noise to sample a counterfactual text from the intervened model. This procedure is described in lines 361–367, but we would greatly appreciate any feedback on areas that might introduce confusion.

Please let us know whether this explanation, alongside the revisions, answer your concerns.

[1] Louis L Thurstone. Psychophysical analysis. The American journal of psychology, 38(3):368–389, 1927.

[2] Michael Oberst and David Sontag. Counterfactual off-policy evaluation with Gumbel-max structural causal models. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.),

[3] John I. Yellott. The relationship between Luce’s choice axiom, Thurstone’s theory of comparative judgment, and the double exponential distribution.

[4] Maddison, Chris J., Daniel Tarlow, and Tom Minka. "A* sampling."

Identifiability

I don't buy the argument in appendix C. There's a bait-and-switch. I completely agree the motivation that there is something unnatural about the inverse cdf model: namely that is not permutation invariant. But you have not proved that all SEMs that have this kind of permutation invariance are of the given form ---- you have instead connected it to Thurstone random variables (Definition C.1). But, to my mind, this assumption is no less contentious than the one you are trying to derive. Why should it be the case that there is a maximization involved at all? Another approach to giving a causal model that is not sensitive to permutations in this way, is to take each $U_i$ to take on values who are functions from the other parent variables to target variables, i.e., use Lubin's "response variables". (I would argue that this SEM is even more natural, as it is the smallest one with certain counterfactual independence.)

Indeed, the authors themselves are aware of this fact. As they write after Theorem C.1,

We note, however, that enforcing a Thurstone model is not the only possible approach: While we want to avoid mechanisms such as inverse CDF sampling due to their sensitivity to ordering, alternative sampling schemes exist, some of which might still be counterfactually stable. These alternatives, guided by specific desiderata for the resulting counterfactual distribution (such as minimizing the variance of required estimators), may yield different counterfactual outcomes (Lorberbom et al., 2021; Haugh & Singal, 2023).

I completely agree. But this begs the question: why even bother with Theorem C.1 then? What exactly is special about the Gumbel mechanism (or, equivalently, the Thurstone model)?

As it happens, I am actually quite sympathetic to the authors' intuitions that there is something special about the Gumbel mechanism. But as it stands, I feel this paper significantly over-promises and under-delivers on the promise of establishing this.

For these reasons, it is imperative that the authors follow through with their choice to drop the word "true" from the title, and otherwise tone down the rhetoric so that it no longer suggests that this is "the right" way to get counterfactuals. As it stands, the paper can only claim a way to get counterfactuals (and perhaps an especially useful or convenient). It seems to me that there is no evaluation of other causal models as a baseline. Question: experimentally speaking, can you see issues with the cdf model (or other permutation-dependent causal models), in comparison to the Gumbel model?

I will revisit my decision during the reviewer discussion period, but at the moment, I still feel that this paper is not ready for publication. It has promise, and the authors clearly have a compelling vision. So it would be a shame to publish this paper prematurely, i.e., before the vision is soundly supported by the science.

FURTHER COMMENTS

• Definition C.1 (specifically, equation 9) does not typecheck; the right hand side is a random variable, not a probability distribution.

• change spacing or add parens in eqn (3)

Here are some options, ordered as I would prefer them:

1. Consider an infinite family $\mathcal M_1, \mathcal M_2, \ldots$ of truncated SEMs. If the generated string has $n$ tokens, then the SEMs {$\{ \mathcal M_i : i < n\}$} in this family that are truncated at length $i < n$ will not generate an EOS token, and all SEMs truncated at $i > n$ will agree on the generated string including the EOS token.

2. Fix a maximum length $M$ of sentence. This is actually quite a reasonable assumption; in practice, we are not interested in sentences that have more than $10^{30}$ words, for instance. I personally do not find think the theoretical contribution of this work is substantially richer for considering sentences of unbounded length. You can always mention that the restriction is artificial, and interventions, etc. can be extended to the infinite case as well.

3. Stick with GSEMs. Put the details (including the definition itself) in the appendix; what's important is just that it allows for infinite variables. State that you're defining interventions in the normal way, and hence not making use of most of the "generalized" part of GSEMs.

Dear Reviewer F4Ae,

We have updated the submission that addresses the discussed concerns (the important changes are now marked in blue). In particular, we highlight the following modifications:

1. We thought deeply about your suggestion about truncation. And, we want to emphasize again how much we appreciate the time you have put into making a detailed technical suggestion. However, we opted to go another way (as mentioned in our past comment) in removing the GSEM definition, as suggested by you in a previous comment. Our revised manuscript offers a definition of a well-founded SEM (generalizing an acyclic SEM) that naturally serves our purposes. We can think abut this definition as just the relevant parts of the GSEM definition without the added overhead. It simplifies our exposition of interventions, but also is sufficient to guarantee a unique solution, for which we cite Halpern and Peters (2021), rather than quoting their definition. We hope this alleviates the concerns.
2. The manuscript now contains significant hedging in the introduction and an improved title. Additionally, we have included a lively discussion about identifiability. While we have sharpened our arguments for why a Gumbel distribution is a useful choice (and a natural first step), we also discuss other options. Indeed, we believe the simplicity of the Gumbel-max parameterization of the softmax paves the road for future inquiry into the right (or at least a good) causal structure behind a language model—we believe this could be a valuable contribution of the added content. We hope this clarifies our objective and makes the contributions more in line with the motivation.
3. We have also rewritten the sampling algorithm in more abstract pseudocode (Alg. 1) and provided a more precise proposition (Prop. 3.1) that supports it. We have also included a detailed proof in the appendix that references the work cited in the previous versions.
4. As suggested, we have included a paragraph explaining the experimental setup in more detail.
5. The final version also addresses some minor concerns, such as defining a function that directly computes the logits defined by a language encoder, $\mathbf{E}$, and $\mathbf{b}$, as well as more careful treatment of random variables and their outcomes.

We hope the modifications address your concerns. We would also like to sincerely thank you for the time and effort invested in your engagement with the work; we believe your feedback resulted in a much better manuscript! We look forward to hearing your final thoughts on the draft.

Dear Reviewer F4Ae,

We hesitate to write again given how much time you have invested in our paper. However, we wanted to ever so gently nudge you to see if our final draft manages to assuage your concerns. Thanks again for making this a very helpful review process for us.

We appreciate the thorough evaluation of our submission. We would like to address the concerns and suggestions you raised.

The reviewer questions the completeness of the noise model, specifically the assumption that the Gumbel noise is independent at each timestep. There is also concern about whether any dependencies in the noise are appropriately captured by the endogenous variables. The independence assumption for the Gumbel noise at each timestep is a key property leveraged by the Gumbel-max trick. The independence does not imply that the noise is uncorrelated across the entire sequence but rather that, conditioned on the observed data, the noise realizations at each timestep are independent given the structural dependencies of the model (i.e., the autoregressive nature of the language model). This is just an expression of the fact that the randomness in sampling is independent between time steps (when we sample tokens from the model, the randomness in choosing token i is independent of the randomness in choosing token j). The dependencies in the sampling process are captured by the deterministic part of the model (the logits computed from the language encoder), which act as endogenous variables in our GSEM formulation. The exogenous noise variables (the Gumbel noise) influence the sampling only through their interaction with these logits. We do not observe any limitation in modeling the noise using a gumbel. It is a common modeling choice. Furthermore, unlike other sampling procedures, like inverse CDF sampling (“roulette wheel”), it is invariant to the arbitrary indexing of the elements in the categorical distribution.

We also ask to emphasize that in our experimental evaluation, we use Algorithm 1 to completely control for all randomness: we make sure that the only thing that changes is the intervention in the model’s parameters. As such, the experiments allow us to isolate, specifically, the influence of the intervention on the output of the model (the tokens sequence). Thank you for pointing out the need for a limitations section. We will add it to the final version of this paper.

We thank the reviewer for their thoughtful and detailed feedback. We appreciate the recognition of our novel framework and the systematic evaluation of intervention techniques. We also value the suggestions provided to improve the clarity and empirical validation of our method. We address the concerns and questions raised as follows:

1. Justification of Proposition 3.1 and Hindsight Gumbel Sampling

The reviewer questions the theoretical foundation of our Hindsight Gumbel Sampling algorithm, specifically challenging the independence of the sampled Gumbel noise and suggesting that the marginal distribution of $U^t$ given an observed sentence might not remain $\text{Gumbel}(0, 1)$. We want to explain the key justification provided in Proposition 3.1.

The central insight of Proposition 3.1 is that the Gumbel-max trick allows us to exactly recover the posterior distribution $P(U \mid \text{observed sentence})$by leveraging the independence property of the Gumbel noise. The Gumbel-max trick ensures that the noise$U^t(w)$for the sampled token$w$(i.e., the observed token) can in fact be modeled as$\text{Gumbel}(0, 1)$. This follows from the key property of the Gumbel distribution, the independence between the max value and the rest, which allows us to sample from the posterior by first sampling a standard gumbel and then sampling the rest of the noise from the truncated distribution.

By conditioning on the observed outcome (i.e., the chosen token), we ensure that $U^t(w_t) + \pi(w_t) \geq U^t(w) + \pi(w)$ for all $w \neq w_t$. This condition precisely defines a truncated Gumbel distribution for the remaining noise variables. Thus, our method correctly samples from the posterior distribution of $U$. We emphasize that this approach is grounded in well-established properties of the Gumbel-max trick, and the independence of the Gumbel noise is key to its validity. We refer to the proof for this property in the appendix.

2. Clarifying the primary contribution and the counterfactual encoder

We appreciate the feedback regarding the presentation of our primary contribution. While the introduction outlines our reframing of language models through causal inference and the novel counterfactual generation method, we will revise the early sections to provide a clearer and more explicit emphasis on the counterfactual generation as a core contribution. We agree that the description of obtaining the counterfactual encoder could be more explicit. In our experiments, we use existing intervention methods to define the counterfactual encoder (e.g., MEMIT, linear steering). We will make this clearer in the revised draft and emphasize that our contribution lies in the generation of true counterfactual strings and the analysis framework, not in proposing new intervention methods.

3. Addressing concerns on sampling variability

The reviewer pointed out that in our evaluation, we sample a single counterfactual, whereas we actually have a distribution over counterfactuals. This observation is valid. However, since our results are negative (i.e., we demonstrate that the counterfactuals tend to deviate significantly from the originals), using a single Monte Carlo sample, averaged over hundreds of different texts, is sufficient to reveal the differences between the originals and the counterfactuals. It is unlikely that we would consistently observe such a difference with just one sample but not with a larger number of samples (N >> 1). Nonetheless, to address this concern empirically, we conducted an additional experiment where we generated 10 counterfactuals per original sentence for the mimic-gender interventions. The results show a high mean cosine similarity of 0.88 between each group of 10 counterfactuals under the E5 model, indicating a strong degree of semantic similarity. Furthermore, we repeat the analysis of the normalized length of the longest prefix between the original and the counterfactuals. The median length changes from 0.281 (for N=1) to 0.265 (for N=10). We will repeat all experiments with N=10 for the final version of this work.

Finally, we appreciate the suggestion to explore counterfactuals of counterfactuals as an additional robustness check. We agree that this could provide deeper insights into the reliability and utility of our method. We will include such experiments in the final version of this work.