Fine-Tuning Pre-trained Language Models for Robust Causal Representation Learning

Although the idea is interesting, the main issue of the paper is that the theoretical basis is not explained in a clear and sound way. Many typographical errors and undefined terms greatly hinder the comprehension of the paper:

1. The transportability concept, mentioned on line 187, is not properly defined and no references to prior work (e.g. [1,2]) are provided
2. Theorem 2 identifies an interventional query $P(Y|do(X)$ based on the causal graph in Figure 1(c), but the variable X is not in that graph
3. In Equation 3, $r_0$ and $r_1$ are not defined
4. There is an inconsistent use of bold and calligraph fonts, for instance: are $\mathbf{X}$ (appears in Assumption 1), $X$ (appears in Figure 1) and $\mathcal{X}$ (appears in the "Motivations and Intuition" paragraph) the same quantities?
5. In Algorithm 1, line 8, the term $\bar{r_i}$ is used, I assume it should be $\bar{x_i}$
6. Again, in Algorithm 1, line 10, $f_\phi$ is used without being defined
7. I have also found many writing errors (e.g. missing words in sentences or wrongly spelled)

[1] "A General Algorithm for Deciding Transportability of Experimental Results" [Bareinboim and Pearl, 2013]

[2] "Transportable Representations for Domain Generalization" [Jalaldoust and Bareinboim, 2024]

Unclear Method Description. The authors use extensive mathematical derivations to analyze and explain the OOD generalization of LMs from a causal inference perspective. Although I appreciate the authors' efforts, I find it unclear how the method is actually implemented, particularly in Section 5. I would appreciate it if the authors could detail the implementation process without relying on math, specifically how a local value is constructed (Section 5.3) and how the model is trained and used for inference (Section 5.4).

(Possibly) Insufficient Evaluation. While the proposed method shows effectiveness in the experiments, the choice of datasets appears relatively simple, and the LM backbone is somewhat limited. The authors use two sentiment classification datasets for evaluation, with bert-base-uncased as the LM backbone.

• Regarding the model, given the significant advancements in existing large language models (e.g., LLaMA3-8b), I am concerned that the proposed method may not remain effective, as such models can easily handle a simple task like sentiment classification and achieve near-perfect performance. It would be helpful if the authors could demonstrate the method's effectiveness on stronger models.
• Regarding the task, I also recommend that the authors evaluate their method on more challenging and commonly-used datasets, such as instruction-following and safety alignment. For example, some large language models like LLaMA-2-7B-chat are known to be over-conservative. During training (RLHF), they are fine-tuned to reject answering user questions that may lead to unsafe outputs. However, this often results in over-conservatism, causing the model to reject many safe requests. Could this be a sign of poor OOD generalization (in terms of safety) in these models? Is it possible to apply the proposed method to these tasks?

I acknowledge that the second suggestion may be challenging to implement, but I encourage the authors to consider tasks beyond simple sentiment classification.

Writing Issues: There are several writing issues in the paper:

• What do $R_1$ and $R_0$ mean? When the authors introduced them in Line 243, they refer to "variations of its (text's) representations," but in Line 261, they claim "r getting its sentence summary $R_1$ and $R_0$."
• Potential typos: Line 336: "obtrain" should be "obtain."
• Possible errors in Table 2: SWA outperforms the proposed method in the 2nd and 3rd columns.
• Some full names and abbreviations are redundant, such as "pretrained language models (PLMs)" in the conclusion section. Since PLM was already defined earlier, it is unnecessary to restate the full term. If this section was revised by ChatGPT or another LLM (which often introduces abbreviations in its revisions as it does not have the full context of previous sections), I recommend that the authors review and remove these repetitions.

However, since I am not an expert in causal inference, I acknowledge that my assessment may be limited. I will adjust my review based on other reviewers' feedback.

• I find the paper difficult to follow
• The developed method is overly complicated, and there might be simpler ways to do it
• The assumptions are very strong, but meanwhile I don't see the clear motivation of some design choices
• There are many typos, sometimes in importance places (e.g. Step 4 in Algo 1) I will elaborate in questions part. I can adjust the points if my questions are addressed.

• The semi-synthetic and real-world experiments appear very similar. Both appear to be based on Amazon and Yelp review datasets, but the injected text of the real-world experiments are based on data source. Exploring datasets suggested in https://arxiv.org/pdf/2311.08648 may strengthen the experiments section.
• More details on the connection to Theorem 4.4 in https://arxiv.org/pdf/2106.04619 would be appreciated.

• In Sec 4.2 theorem 1, the authors assume that the causal latent variable C between two environments (pre-trained and fine-tuned models) stays the same. However, this assumption seems a kind of logical leap, as there is no supporting experimental analysis or general consensus, and at the same time, practically not trivial. In other words, this assumption needs that practical solution (good fine-tuned model) is already fulfilled; if there exists a causally robust language model on each downstream task, why do we need OOD generalization of the pre-trained model?
• Experiments are monotonic, so it is hard to comprehensively evaluate whether the suggested method increases OOD generalization.
  • The available experimental results in the main body seem to repeat the same experimental design over just two datasets (though there are plenty of available classification benchmarks), which is insufficient for holistic understanding. Could you provide additional experiments on other dataset in diverse domain?
  • There is no experimental analysis other than performance, so it is hard to understand how the designed algorithm actually works regarding causality. Could you provide more in-depth analysis not only performance?