Exploring How LLMs Capture and Represent Domain-Specific Knowledge

We sincerely thank the reviewer for their insightful comments and constructive feedback. Below, we address each point raised.

Dataset and domain characterization: The decision to select 30 out of the 57 domains from the MMLU dataset was driven by the need to focus on specific domains where the required skills for LLMs differ significantly, providing a clearer spectrum for analyzing contrastive behavior. This approach was inspired by the AdaptLLM [1] paper, which demonstrated diverse skill requirements across domains such as Biomedicine, Finance, and Law. Our comparative analysis of models based on the Llama2 backbone, included in Figure 7, builds on these models.

Additionally, we included the mathematical domain due to their large sample size within the MMLU dataset and the abundance of publicly available fine-tuned checkpoints. This allows for broader comparisons and a more comprehensive evaluation. Ultimately, 30 subdomains qualified under these domain categories, as detailed in the GitHub repository of the dataset and listed in Appendix A.1. The remaining subcategories, such as those categorized under miscellaneous or global facts, were excluded to avoid ambiguity and ensure clearer, more interpretable results.

It is important to note that only the MMLU dataset underwent this filtering process; the other eight datasets were utilized in their entirety without filtering any samples. We thank the reviewer for their valuable feedback and will enhance Appendix A.1 to include a detailed list of the selected subdomains, along with additional information about their content.

Next steps/ Research directions: We envision that the findings from analyzing domain-specific hidden state patterns can have several practical applications:

Steering Model Behavior: By modifying hidden states in real time along “well-known” trajectories, it becomes possible to guide the model’s outputs toward specific behaviors or styles. For example, steering vectors can be employed to influence the generation process within a specific domain, enhancing the model’s adaptability.

Data/Training Efficiency: Hidden state traces can serve as a proxy for identifying features that drive faster model convergence, facilitating more efficient use of training data. This concept aligns with ideas like the “Lottery Tickets Dataset,” enabling optimization of data utilization and reducing training costs.

Sparse Model Optimization: In our recent experiments with the Phi-3.5-MoE model, we observed that sparse models exhibit predictable activation patterns in the early layers of their hidden states. These patterns could potentially allow for early predictions of which experts will be activated, paving the way for dynamic load balancing and smarter compute allocation.

These examples represent just a few promising avenues for further exploration. While they extend beyond the scope of this paper, they highlight significant opportunities for future research. We will update the final discussion to better articulate these potential directions and provide additional clarity.

Regarding appendix A.3, A.4 and A.5: Due to space constraints, we decided to relocate this content to the appendix to ensure sufficient focus on highlighting the main contributions of this work. To maintain accessibility, we have added footnotes in the main text with clear pointers to the corresponding content in the appendix. We will ensure this is further emphasized in the final version for improved clarity and navigation.

Regarding hyperparameter selection for baseline method: For hyperparameter selection across all baseline methods, we ensured the best achievable performance for each respective model while using the same training samples across all experiments. This approach allowed us to maintain a fair comparison under consistent conditions, ensuring an “apples-to-apples” evaluation framework.

We sincerely thank the reviewer for their insightful comments and constructive feedback. Below, we address each point raised.

Insufficient theoretical foundation: We acknowledge that our study primarily focuses on empirical observations rather than presenting a fully developed theoretical framework. The central aim of this work is to experimentally investigate hidden state patterns to identify trends that suggest domain-specific processing by utilizing information from all layers.

Regarding the "latent domain-related trajectories" introduced in the paper, we agree that their correlation with domain representation would benefit from more robust mathematical grounding. Our intention was to showcase consistent patterns across diverse domains—such as medical, legal, and mathematical—when analyzed across various architectures and datasets. These patterns, while not constituting definitive theoretical proof, provide preliminary evidence of the influence of domain knowledge on model intermediate activations on the prefill phase.

We also appreciate the reviewer's suggestion for more rigorous statistical validation of the variance computations in equations (1) and (2). While our current analysis provides a measure of variability in activations linked to domain-specific tasks, we acknowledge the need for further justification. To address this, we are incorporating entropy across activations as an additional validation metric in our experiments for the final version. Entropy has been widely recognized as a robust measure for quantifying uncertainty and variability in model behavior, offering a complementary perspective to variance.

Model selection in Section 4: Our decision to focus on smaller language models (below 7B parameters) was primarily driven by resource constraints, as experimenting with larger models demands significantly more computational resources. We acknowledge this limitation and have explicitly noted in Section 6 that the applicability of our approach to larger models remains an avenue for future investigation.

We recently extended our experiments to include the Phi-3.5-MoE instruction-tuned model (41B parameters). These experiments revealed that sparse models exhibit distinct separation of traces from the early layers, offering exciting possibilities on other research areas such as dynamic expert allocation. These findings, which expand the scope of our conclusions, will be included in the final version appendix to provide further context.

We respectfully disagree with the suggestion that our model selection seems arbitrary. Our approach was deliberately designed to validate our findings across a diverse range of widely-used open-source model architectures, training methodologies, and parameter sizes. This variety ensures that our conclusions are both robust and derived from varied experimental setups, enhancing the generalizability of our observations. Should the reviewer have any specific concerns regarding our model selection or require further clarification on any aspect, we would greatly appreciate the opportunity to address them.

Domain categorization lacks systematic justification: The domain categorization for each dataset was determined based on the nature of the questions, with all datasets belonging to a single domain except for MMLU, whose categories are detailed in Appendix A.1, as explained in Section 4. We are uncertain about the specific aspects of our approach to domain categorization that you find unclear or lacking. Could you kindly elaborate on this point so that we may address your concerns more effectively?

Prompt consistency analysis: In our experiments, we tested four different prompt variations per domain (see Tables 1, 4, and 5), while maintaining the core intent of each question. In some cases, we deliberately omitted context to create more challenging scenarios for the model. The results for these cases are shown in Figure 3, where we observe minimal variance across the last 10 layers of the baseline model. Additionally, we include variations across the fine-tuned models in Figure 6. In all cases, we observe that the model tends to group queries based on domain similarity.

As noted in Table 2, each dataset has its own prompt variation (for example, although GSM8K and MATH are from the same domain, each dataset uses a specific variation, as detailed in Table 1). We believe we have sufficiently explored prompt variation in our setup, demonstrating robustness across different scenarios. However, we welcome any further suggestions on how we might enhance this analysis or explore additional prompt variations to strengthen our findings.

We sincerely thank the reviewer for their insightful comments and constructive feedback. Below, we address each point raised.

Regarding clarity of the paper: We appreciate the reviewer’s feedback regarding the clarity and engagement of the writing. To address this, we will revise the manuscript to improve readability and flow. Specifically, we will make sure to highlight the potential extensions of this work in practice for the final version of the paper.

Rationale behind domain sensitivity: The observed hidden state patterns suggest domain-specific understanding primarily through experimental evidence rather than a direct assertion. As outlined in the paper, we explored hidden states across multiple architectures and datasets to observe whether consistent patterns emerged. This approach is rooted in the idea that domain-specific understanding should manifest in the way a model activates and processes information across layers when presented with context-specific inputs [1][2][3]

The patterns we observe are not inherently conclusive but show distinct trends that suggest domain sensitivity. For example, across various domains like medical, legal, and mathematical tasks, we saw that specific layers (second half of layers) tended to activate more strongly in response to relevant tokens or concepts tied to each domain. This observation was consistent across multiple model architectures, which strengthens the argument that these patterns are not simply artifacts of model design but rather indicative of domain-related reasoning processes.

However, we acknowledge that these patterns do not prove a deep, inherent understanding of the domain in a traditional sense. Rather, they reflect domain-related processing and response patterns that differentiate between various types of information (e.g., factual, procedural, argumentative) across domains. The variation in these activations across datasets further supports the notion that the models are exhibiting domain-specific behavior, even if it remains somewhat superficial in terms of understanding. Therefore, our experimental methodology—testing across different architectures, finetuned models and using a wide variety of datasets—aims to provide evidence for domain sensitivity in hidden states, even as the exact nature of this sensitivity remains a subject of further investigation.

Further applications of this work: The insights from analyzing domain-specific hidden state patterns can be applied in several practical ways:

• Steering Model Behavior: Modifying hidden states in real-time, towards “well-known” trajectories we can guide the model’s output toward specific behaviors or styles, such as using steering vectors to influence in specific domain generation process.
• Data/Training Efficiency: The hidden state traces can act as a proxy to identify which features drive faster convergence, enabling more efficient data use (e.g., the “Lottery Tickets Dataset”).
• Sparse Model Optimization: In our recent experiments on Phi-3.5-moe we have found that Sparse models exhibit predictable activation patterns in their hidden states of early layers. This insight could potentially allow for early prediction of which experts will be activated, enabling dynamic load balancing and smarter compute allocation.

We leave these research angles for future directions as extensions. We thank again to the reviewer for their detailed feedback, and we will include the suggestions into the final version of the paper.

[1] Can LLMs Infer Domain Knowledge from Code Exemplars? A Preliminary Study.

[2] Linearity of relation decoding in transformer language models.

[3] Inspecting and editing knowledge representations in language models.

We sincerely thank the reviewer for their insightful comments and constructive feedback. Below, we address each point raised.

DeBERTa's lack of separation stems from model size: We appreciate your observation regarding size differences influencing domain separation in DeBERTa. Additional experiments on autoregressive models of similar size (GPT-2, GPT-Neo, OPT at 125M parameters) confirm that separation patterns persist in smaller autoregressive models, though less pronounced. These findings will be included in the appendix, along with updates to Section 5.1 to clarify the model size’s role.

Overlap between medical and math traces: We agree with your observation of a significant overlap across these domains. In the detailed analysis of our hypothesis in Appendix A.5, we explain that this overlap may stem from the shared structural reasoning processes required in these fields. In contrast, we observe a smaller overlap in the fields of laws and humanities, where the reasoning process relies more heavily on persuasive argumentation.

Routing results and dataset performance: The primary goal of using hidden state activations for routing is to demonstrate their potential in selecting a model that best aligns with the unique characteristics of the input query. This is particularly valuable in unsupervised routing for open-ended and closed generative tasks, as the routing is guided by hidden state patterns rather than explicit sample labels.

To train the router, we utilized hidden states from the Phi-3-mini-128k model applied to 4,000 random samples from the MMLU dataset (Base Pool). During inference, we evaluated the router on 1,000 unseen samples from each dataset listed in Table 2, including GSM8K, MATH, MEDMCQA, USMLE, and CASEHOLD. The router learned to classify queries into four domains—maths, biomedical, law, and humanities—and used this classification to route queries to the corresponding fine-tuned models for each domain. Additional clarifications:

• the performance metrics in Table 2 are based on these 1,000 samples per dataset.
• Standard errors and dataset sizes will be included in the final version.

Regarding cases where, e.g., a mathematical query may benefit from a medical-domain model: in our experiments, we observed that the performance differences between domain-specific models can be small for certain tasks, suggesting that reasoning patterns or shared latent representations may overlap across domains. This overlap enables the router to identify subtle but meaningful connections that influence routing decisions effectively. Regarding the observed lower performance on the MATH and GSM8k datasets, this is attributable to the nature of the tasks. These datasets primarily contain open-ended, complex questions, which are inherently more challenging than constrained formats like multiple-choice questions. This reinforces the value of our approach in navigating these complex scenarios and highlights the potential of hidden state patterns in improving task-specific model selection.

Equation 1: While layer normalization ensures a mean of zero for each token across its feature dimensions during the normalization step, the activations per layer (A) represent the final layer outputs, which include additional transformations (residual connections and learned bias terms applied after normalization); these transformations alter the mean, making it non-zero when aggregated across the batch and feature dimensions. Additionally, when aggregating across the batch and feature dimensions, variability in the inputs and context further contributes to deviations from zero.

Table 2: The hidden states in Table 2 are derived from Base Pool samples obtained with Phi-3-mini-128k model (not finetuned) to maintain consistency with the traces across its finetuned counterparts. As discussed in Section 5.1 and Appendix A.4, we show that these traces persist across the different finetuned versions, ensuring consistent interpretability across different versions of the same model.

(L379-L388) Regarding fine-tuned models used in the router: to clarify, the checkpoints from Hugging Face used in our experiments were not fine-tuned by us (due to resource constraints). Instead, we selected stable, publicly available checkpoints tailored to different domains, all based on the same small base model used in our experiments (Phi-3-mini-128k). Specifically, we selected high-quality models fine-tuned on math, medical, and emotional domains. By "high-quality” we mean checkpoints with a clear training process and strong performance on unseen datasets. These are the four models (3 finetuned + pretrained) evaluated in step 3 of section 5.3. After evaluating their performance, we identified the strongest checkpoints, as detailed in Appendix A.4.