Enhancing Robustness of
Foundation Model Representations
under Provenance-related Distribution Shifts

In the simulation of different degrees of distribution shifts by provenance, 4 parameters are required to set up the framework: $P_{train}(y=1|z=0),P_{train}(y=1|z=1),P_{train}(z=1),\alpha_{test}$. The first three were sampled from 0 to 1 evenly in linear space, with a step size of 0.05 or 0.1, depending on scenarios.

Since $\alpha$ represents the ratio for positive rates from two sources and $\alpha=1$ means same prevalence rates, we sampled $\alpha_{test}$ in a reciprocal manner while centering around 1.0. One such example is $\{...,1/8,1/4,1/2,1,2,4,8,...\}$. This can ensure a uniform distribution of $\alpha_{test}$ in the log space, as shown in Figure A1.

To break down $\alpha_{test}$ into detailed samplings of its two subpopulations, Figure A2 shows the theoretical sampling for the joint distribution of $P_{test}(y=1|z=0)$ and $P_{test}(y=1|z=1)$. It is worth noting that not all sampling configurations from Figure A2 can lead to a valid training/testing set, given how many positive/negative samples from both sources required for that setting.

To generate embeddings for sentences, we applied the Sentence-BERT model [18]. We adopted the publicly available version of the model all-MiniLM-L6-v2 from the HuggingFace repository ²²2https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2. Sentence-BERT produces a single embedding for a given unit of text. The resulting sentence embeddings serve as predictor variables in the regression model.

In this work, we extracted sentence embeddings from all three versions of the Llama 2 model (without additional training for dialog). Officially, they are named after how many parameters (roughly) each of them contains: Llama-7B, Llama-13B, and Llama-70B. Embedding dimensions from Llama-7B, 13B, and 70B model are 4096, 5120, 8192, respectively. The first two can be fit into one A100 GPU, while for the largest 70B version, we applied 8-bit quantization [30] before loading it into GPU. For each token, we only kept the outputs from the very last attention layer. From embeddings, we took average as the pooling function (as oppose to use the embedding from the last non-padding token) over each dimension for all tokens in one sentence. This was demonstrated with better performance in the preliminary work.

In this section, we present baseline results on additional results on Llama 2-13B and quantized Llama 2-70B, for the SHAC dataset (Figure A3, coefficients presented in Table A3) and full results for the Hate Speech Detection dataset (Figure A4, coefficients presented in Table A4). Besides of results on $\alpha_{train}=0.33$, Figure A5 and Figure A6 show additional results on $\alpha_{train}=3$, as the reciprocal according to Simulation Configurations in Appendix A. It is noted that the value pair of 0.33 and 0.3 for $\alpha_{train}$ is randomly selected as an example and our interpretation holds true for other $\alpha_{test}$ values (not presented) with differences in degree of the effects. Different $\alpha_{test}$ ranges are shown for results on two datasets because empirical sampling may invalidate some of theoretical combinations.

For unadjusted models, embeddings from foundation models - Llama 2 models in particular - provide best innate robustness to confounding shifts. With respect to this robustness, Sentence-BERT lies at a similar level to the baseline binary unigrams. As for Llama 2 specifically, a larger model in size does not confer more robustness, as shown in our results on the two datasets where Llama 2-7B can produce lowest absolute coefficients in some scenarios ($C_{y}$ levels).

After applying Backdoor Adjustment to the logistic regression models, the absolute values of coefficients for all models drop, indicating better model robustness to provenance-related distribution shifts. Among these, the baseline binary unigrams exhibit the most significant improvements after the adjustment, with the lowest absolute coefficients. However, for results on both datasets, the AUPRC measures from this baseline model are typically lower and deteriorate rapidly in comparison with those using foundation model representations. This performance drop is more significant with the Hate Speech dataset (Figure A4).

When reviewing results for $\alpha_{train}=3$ in Figure A5 and Figure A6, we observe similar trends. Due to the sampling strategy, only some $C_{y}$ groups can be matched as in the case where $\alpha_{train}=0.33$. One apparent difference is that the slope of the lines is now mostly positive, as opposed to the negative slopes with $\alpha_{train}=0.33$. This indicates the training set can influence the model’s baseline performance, in that when the testing set is more different from the training set (in terms of $\alpha$ values in our case), the performance usually drops. This matches the far right area when $\alpha_{train}=0.33$ and far left area for $\alpha_{train}=3$.

Overall, the empirical experiments on both datasets show consistent results. The logistic regression model, without any adjustment, is very sensitive to distribution shift, as evident in the higher absolute values of the coefficients of lines fit to the performance curves. However, different embedding methods show various degrees of robustness. When no adjustment is made, models trained on Llama 2 embeddings show best performance, including the smaller 7B or 13B versions, especially when the dataset is small. As with the SHAC dataset, out-of-the-box usage of embeddings from foundation models (e.g., Llama 2-7B) already shows some robustness to confounding shifts. However, without adjustment these representations do not provide significantly better robustness over the baseline unigram model (as in a larger Hate Speech Detection dataset). All these results suggest a need for adjustment to logistic regression models in the context of provenance-related confounding shift, irrespective of the choice of text representation technique.

Backdoor adjustment, when applied to the logistic regression, can significantly improve robustness to these provenance-related shifts. This holds true for different embedding methods on the two datasets evaluated. Among them, binary unigrams benefits most from the adjustment, as shown in Figure 1(a) when $C_{y}=0.48$, especially for the larger Hate Speech Detection dataset in Figure 4(a). However, this comes with a loss in performance of AUPRC, which in comparison is not significant for the SHAC dataset with very high baselines.

One reason that Backdoor Adjustment works relatively well on unigrams could be that foundation models generate highly contextual embeddings. In comparison, binary unigrams only store information at the unigram level from any corpus. These two ways of representing text capture different information, which may further affect the distribution of representations. Moreover, it should be noted that the embeddings from different models have different dimensions, which subsequently affect the size of a regression model ($\boldsymbol{\beta}$). As such, it could be argued that models trained on the resulting embeddings are not strictly comparable. A potential way of controlling for this would be to apply dimension reduction methods, to enforce the constraint that all inputs to the regression model must have the same dimension. However, this will inevitably cause loss of information, with effects that are as yet unclear. There is another hyper-parameter, $v$, applied to the one-hot matrix for Backdoor Adjustment. This serves as an added regulation term in regression models, and may have different effects with different numbers of parameters. In our work, $v$ was set to 10 for all experiments. The optimal $v$ value for models of different sizes remains to be determined. Beyond our empirical experiments, formal testing of these assumptions is left to future work.

In this work, we extended an evaluation framework for provenance-related confounding shift and examined one strategy for mitigating its effect on model stability. Several constraints were enforced during the development of the framework. They were originally set to isolate the effects of changes in site-specific class prevalence, but at the same time limit the scope of testing. Further work is also required to expand the framework to multiple classes and sources, other than binary cases discussed in this work. Moreover, it is left to future work to expand the framework to include other confounders, as in the case of the CivilComments dataset where a combination of 8 dimensions (e..g, gender, religion) collectively serve as the domain label [31].

The unprecedented size of foundation models (Llama 2 especially) limited our ability to utilize it fully in the context of available resources, for example by preventing us from fine-tuning it end-to-end. The path we took in this work is a computationally lightweight approach and only serves as a first step into exploring Llama 2’s potential. It remains to be determined whether such foundation models can be more robust when tuned using a parameter-efficient approach, or even fully fine-tuned.

Our primary focus has been on the application of foundation models to classification tasks in multi-institutional datasets within the biomedical domain, though several observations from our biomedical set align well with results from the general domain dataset. Further validation is required to determine how these findings generalize to different scenarios, including those with larger training samples and different model architectures, as well as whether different biomedical subdomains will favor different strategies.