Transformer Copilot: Learning from The Mistake Log in LLM Fine-tuning

We sincerely thank the reviewer for your extensive and thoughtful feedback! We will address each of your remaining concerns below.

---

[W.1] Theoretical Justification for $\lambda$

Thank you for this insightful question! Below, we justify both theoretically and experimentally.

"Is it possible to extend these to consider the vector as a whole?"

Theoretically: Yes, Theorem 1 can be extended to establish that there exists a constant $\lambda$ such that the theoretical guarantee holds for all $K$ dimensions (where $K = |V|$). Specifically, recall that for each dimension $k \in [K]$, there exists a threshold $\lambda_0^k$ such that the guarantee holds for that dimension when $0 < \lambda < \lambda_0^k$, where $\lambda_0^k = \frac{2\sqrt{\epsilon_P^2+\sigma_P^2}\big(\sqrt{\epsilon_P^2+\sigma_P^2}-\epsilon_C\big)}{\big(\sqrt{\epsilon_P^2+\sigma_P^2}-\epsilon_C\big)^2+\sigma_C^2}$. Let $\widehat{\lambda}_0 = \min\{\lambda_0^k: k \in [K]\}$. Then, the guarantee holds uniformly across all $K$ dimensions whenever $0 < \lambda < \widehat{\lambda}_0$.

Empirically: It is not always necessary for $\lambda$ to fall within this range for all dimensions in practice. As the reviewer pointed out, it is often sufficient to choose a $\lambda$ that satisfies $0 < \lambda < \lambda_0^k$ for the majority of dimensions $k$, which typically leads to good empirical performance (as demonstrated in Tables 1 and 2 of our paper).

---

[W.2] More Quantitative Analysis on Copilot Corrections

Thank you for the helpful suggestion. In response, we include additional metrics in our error-correction analysis to quantify how often the Copilot correctly or incorrectly modifies the final token predictions.

We include the following new metrics:

• Correction Rate: The percentage of test examples where the Copilot changes the final token prediction of the Pilot.
• Correction Precision: Of those examples where the Copilot changes the token prediction, the percentage where the token correction leads to the correct final answers.
• Correction Error Rate: Of the changed examples, the percentage where the Copilot’s correction leads to incorrect final results.

We randomly sampled 1,000 examples from the test sets of PIQA, BoolQ, AQuA, GSM8K, and LastFM to ensure broad coverage across different tasks conducted in our paper.

Results:

The results show that T-Copilot achieves much higher accuracy in correcting the Pilot’s logits predictions, highlighting its robustness and effectiveness in error-correction.

---

[Q.1] Backbone Choice for Copilot Model

"Is there any reason the Copilot model should be of the same backbone as the Pilot model? Could it be possible that training a different backbone from scratch, not fine-tuning, may lead to better results?"

Thank you for this thoughtful question. In short, T-Copilot does not restrict the Copilot to using the same backbone architecture as the Pilot; however, using the same backbone is more practical and simplifies implementation.

For a more detailed response below:

1. We first provide the key conceptual and practical reasons for adopting the same backbone as the Pilot;
2. We then discuss potential extensions of the Copilot to alternative backbones, as suggested by the reviewer.

1. Justification for using the same backbones.

• Unified Vocabulary Space: Using the same backbone ensures that both the Pilot and Copilot share the same tokenization scheme and vocabulary space, enabling the Copilot's consistent token-level logits correction. Using different backbones might require additional vocabulary projection or alignment heads, which complicates the overall integration.
• Compatibible Architecture and Residual Alignment: The Copilot receives the Pilot’s hidden states and residual errors as inputs. Sharing the same backbone ensures representation compatibility, simplifying attention alignment, logits fusion, and decoding dynamics. A radically different backbone may potentially introduce misaligned inductive biases.
• Parameter Reuse and Training Efficiency: Initializing the Copilot from the same backbone allows efficient reuse of pretrained knowledge and ensures stable training. However, training a distinct model from scratch may require substantially more data and computational resources to achieve comparable performance.

2. Extension to different backbones. We completely agree with the reviewer that T-Copilot can be smoothly extended to train a Copilot with different backbones from scratch (e.g., pretrain a Copilot model on residual errors across various tasks), potentially offering stronger error-correction capabilities.

Achieving such (pre-)training scale would typically require substantially more data and computational resources than the current scope of our experiments allows. We appreciate the reviewer’s suggestion and plan to explore this direction in our future work.

---

[Q.2] Have other pooling methods (besides mean pooling) been tested?

Thank you for the question. Yes, we have explored other pooling methods during our early stage method design, including max pooling and sum pooling. We report the results below:

We adopted mean pooling over alternatives because:

• In our preliminary experiments shown above, the mean pooling empirically provided the best balance of stability and efficiency.
• Prior works [1,2,3] suggest that mean pooling is the most common and widely adopted approach for processing transformer hidden-state representations.

References:
[1] Sentence Embeddings using Siamese BERT-Networks.
[2] Comparative Analysis of Pooling Mechanisms in LLMs: A Sentiment Analysis Perspective.
[3] A Brief Overview of Universal Sentence Representation Methods: A Linguistic View.

---

[Q.3] Impact of $\lambda > 1$

"What would be the effect of increasing $\lambda > 1?$"

We thank the reviewer for the detailed question on $\lambda$.

• Intuitively, as shown in Eq.(8), the Copilot primarily serves as an "assistant" to help the Pilot rectify its logits during generation. Therefore, the Copilot’s output should not be weighted more heavily than the Pilot’s (i.e., $\lambda \leq 1$).

• Empirically, we follow our ablation study of $\lambda$ in Appendix G.6 and evaluate the downstream performance when $\lambda > 1$ on LLaMA-3-8B + T-Copilot-1B:

  $\lambda$	HellaS.	GSM8K	SVAMP
  1.0	93.5	66.1	75.4
  1.2	94.2	65.5	74.7
  1.5	93.8	65.8	74.2

  The results show that increasing $\lambda > 1$ does not consistently yield additional gains on tasks such as GSM8K and SVAMP. Therefore, we consider $\lambda = 1$ to be the optimal setting for T-Copilot during generation.

Dear Reviewer sdqu,

Thank you for your positive rating and thoughtful review of our paper! We are pleased that our rebuttal has addressed your concerns. We will incorporate the results and analyses from the rebuttal to further strengthen our paper.

Warm Regards,

The Authors

Thank you for your positive comments and insightful feedback on our paper. We will address each of your remaining concerns below.

---

[W.1] Assumption in Main Theorem 4.1

We thank the reviewer for the insightful question regarding the assumption of Theorem 4.1. We would like to clarify and justify this assumption from both theoretical and empirical perspectives.

1. In theory:
While the assumption $\epsilon_C^2<\epsilon_P^2+\sigma_P^2$ might look strong at first glance, upon closer inspection, it is actually a mild requirement regarding the Copilot’s learning capability.

• As elaborated in Remark 4.2, the assumption implies that the Copilot can tolerate a larger bias (i.e., weaker learning capability) than the Pilot, i.e., $\epsilon_P^2<\epsilon_C^2<\epsilon_P^2+\sigma_P^2$.

2. In practice:
We implement the Copilot under the same decoder backbone architecture as the Pilot, and it is of relatively smaller size.

• We believe it is reasonable to expect the two models to have similar learning capabilities or biases, since the upper bound ensures the Copilot is actually learned during the SFT as the Pilot does.
• We empirically measure on three held-out tasks (GSM8K, SVAMP, and PIQA). In all cases, the Copilot’s combined MSE is 15–20% lower than the Pilot’s (i.e., $\epsilon_C^2 <\epsilon_P^2+\sigma_P^2$), thereby empirically satisfying the assumption condition.

---

[W.2] Ensemble Learning Baseline.

Thank you for the insightful suggestion to compare with ensemble learning variants. In response:

1. We first highlight the key differences between T-Copilot and ensemble learning methods.
2. We then provide an empirical comparison with the ensemble learning variant to demonstrate T-Copilot's error correction capabilities.

1. Difference with Ensemble Learning.
Unlike ensemble learning methods that aggregate independent model prediction outputs, T-Copilot leverages the Copilot trained to explicitly correct the Pilot's errors through the Mistake Log. Instead of generating next-token predictions directly, the Copilot learns to predict the discrepancy between the Pilot’s output and the ground truth, and rectifies the Pilot’s logits accordingly.

2. Empirical Comparison.
To better demonstrate the benefit of T-Copilot's error-correction capabilities, beyond the simple ensemble of two models, we follow the reviewer's suggestion and conduct a new comparison experiment below.

Experiment Setups: We utilize the same training and implementation setups from Section 5 of our paper, and compare with the following variant:

• Pilot-Copilot-Ensemble: Following works [1,2], we train Pilot and Copilot models independently and ensemble at inference by averaging their per‑token logits distribution to select each next token.

Evaluation Results:

• From these results, T-Copilot consistently outperforms the simple ensemble learning variant across different tasks.
• This confirms that T‑Copilot’s improvements are not merely due to “the integration of two models,” but stem from its error‑aware rectification mechanism.

We appreciate the reviewer’s suggestion and will include a thorough discussion and comparison with the ensemble learning variant in our paper, as shown above, to better highlight the advantages of T-Copilot.

References:
[1] Ensemble Learning for Heterogeneous Large Language Models with Deep Parallel Collaboration.
[2] Harnessing Multiple Large Language Models: A Survey on LLM Ensemble.

Dear Reviewer LvNr,

Thank you for your positive rating of our submission! We are glad that our clarifications addressed all of your questions. If you have any remaining concerns, we would be happy to discuss them further. We sincerely appreciate your time and valuable feedback!

Warm regards,

The Authors

We sincerely thank the reviewer for the detailed and insightful feedback. We address each of your remaining concerns below.

---

[Weakness] Ablations on Alternating Self/Cross-attention in the Copilot.

Thank you for acknowledging our experimental efforts. We provide additional ablations on T-Copilot's design choices. For the design choices of Copilot's modified attention, we conduct two detailed ablations.

1. Pooling Methods of Copilot's Attention Input.

We explored other pooling methods, including max pooling and sum pooling. The results are reported below:

From the results, the mean pooling empirically provided the best balance of stability and efficiency.

2. Insertion Pattern of Copilot's Attention.

We also want to highlight that, in Appendix G.6, we investigated four different insertion patterns for Copilot's modified attention mechanism. The full results are reported in Table 18 of our paper.

References:
[1] Comparative Analysis of Pooling Mechanisms in LLMs: A Sentiment Analysis Perspective.
[2] A Brief Overview of Universal Sentence Representation Methods: A Linguistic View.

---

[Weakness & Question] Ablations on RMSE Loss.

"Can you provide intuition ... for error prediction?"

We thank the reviewer for inquiring about T-Copilot's loss types. Below, we provide conceptual rationale and empirical ablations on RMSE loss for T-Copilot.

1. Conceptual Rationale

Recall that the objective of the Copilot is to predict the distribution discrepancy, defined as $\ell\_t(p\_{t,i}, \hat{p}_{t,i}) = p\_{t,i} - \hat{p}\_{t,i}$ in our paper. Note that $\ell\_t(p\_{t,i}, \hat{p}\_{t,i})$ is not a valid distribution, since Sum$(\ell\_t(p\_{t,i}, \hat{p}\_{t,i})) \neq 1$. Therefore, it is natural to formulate this task as a regression problem, for which the RMSE loss is commonly adopted.

In contrast, CE and KL-divergence are for distribution fitting and are not directly applicable here unless we manipulate $\ell\_t(p\_{t,i}, \hat{p}\_{t,i})$ with softmax to resemble a valid distribution. In this case, the loss is not directly optimized on the original discrepancy (i.e., $p\_{t,i} - \hat{p}\_{t,i}$), resulting information loss.

2. Empirical Ablations

In practice, we conducted experiments with all three loss types, and RMSE consistently achieved the best performance.

---

[Weakness & Question] Memory Overhead of Mistake Log.

"How do you manage... Could this limit scalability?"

Thank you for asking the memory storage of Mistake Log.

First, we would like to elaborate on how our current method and implementation are designed to handle the storage of the Mistake Log:

1. Maintenance of the Mistake Log Buffer.
We adopt a buffer mechanism for Mistake Log in our implementation to reduce memory usage in practice:

• All logged tensors are detached from the GPU and moved to a fixed‑size CPU buffer (Lines 894-895); GPU memory is therefore almost unaffected.
• We utilize a fixed‑size buffer that only retains the 128 training rounds in Mistake Log (Line 896).
• Experiments show that adding the buffer reduces the overall storage by roughly 92% while maintaining a similar performance.

2. Mean Pooling on Pilot's Forward Pass.
In Eq. (6), we perform mean pooling across all layers' hidden state representations before feeding into the Copilot, meaning:

• During the Mistake Log correction process, we only need to store the mean pooled vectors rather than the full hidden states across all layers, which significantly reduces vector storage.
• During Pilot forward pass, the hidden states have already computed within transformer's residual stream. We can directly access these hidden representations without allocating extra space for storage or backpropagation gradients.

Second, to further improve the scalability of our method and adapt to much larger models and data, we can slightly modify our design:

1. Selective Storage: We can design a Mistake Log Gate to filter and collect hidden states information only when typical error patterns (mistakes) occur.

2. Lightweight Mistake Log: Instead of curating Mistake Logs over the entire training process, we directly use the Pilot’s information from the current round's forward pass as input to the Copilot.

---

[Weakness] Difference with Existing Ensemble or Multitask Architectures.

Thank you for recognizing T‑Copilot’s architectural novelty and its relation to related methods. We highlight two key distinctions:

• Error‑Aware Supervision vs. Output Averaging.
  Traditional ensembles combine each model’s logits predictions post‑hoc, without specialized training and inference paradigms for error-correction. Our Copilot can explicitly identify and correct the Pilot’s errors.
• Discrepancy Modeling vs. Direct Generation.
  Instead of generating next-token predictions, the Copilot learns to predict the discrepancy between the Pilot’s output and the ground truth and rectifies the Pilot’s logits.

---

[Question] Comparison with Self-Refine and Refinanced Self-Training.

Thank you for your suggestions. Below, we provide a comparison between T-Copilot and Self-refine [1] for better demonstration.

Setups: We compare Qwen2.5-7B+T-Copilot-3B with Self-Refine on GSM8K and MBPP. For Self-Refine, we follow its original protocol with 3 feedback-refine iterations.

Additional Analyses: Self-refine [1] and Reinforced Self-Training [2] rely on external feedback and offline RL, whereas T-Copilot leverages the model’s internal learning signals during SFT.

We sincerely thank the reviewer for the suggestion and will include the experiments above and additional comparisons in our paper revision.

References:
[1] Self-refine: Iterative refinement with self-feedback.
[2] Reinforced self-training (rest) for language modeling.

---

[Question] Generalize to Summarization and Code Generation.

Thank you for the question. Below, we conduct new experiments on code generation and text summarization tasks.

• Code Generation
  • Training: We randomly select 5k examples from CodeContests [1] as the training set.
  • Evaluation: We evaluate the standard test set of MBPP [2] and LiveCodeBench [3] version-2 (511 problems). We report the pass@1 accuracy.
• Text Summarization
  • Training: We randomly sample 10k training examples from the CNN-DailyMail [4] and provide a training split.
  • Evaluation: We evaluate on the standard test set of CNN-DailyMail and report the ROUGE-1 F1-score.

Across all three tasks, T‑Copilot improves performance over the base Pilot, demonstrating the generalizability to diverse generation tasks.

References:
[1] Competition-level code generation with alphacode.
[2] Program synthesis with large language models.
[3] Livecodebench: Holistic and contamination-free evaluation of large language models for code.
[4] Get to the point: Summarization with pointer-generator networks.

---

[Question] Training Stability and Hyperparameter Sensitivity.

Did you observe ... different objectives? How sensitive ... rates?

Thank you for the insightful questions regarding training stability. In our experiments, we did not observe notable instability when jointly training the Pilot and Copilot models:

• Jointly Trained but Separately Optimized: In our training paradigm, Copilot receives inputs via the Mistake Log (i.e., detached information from Pilot’s forward pass). This careful design ensures that Copilot learns from Pilot’s Mistake Log without perturbing Pilot’s weights, thereby preventing interference between the two models and avoiding potential instability or bias.
• Adoption of Standard Optimizers and Schedules: We use AdamW with a cosine schedule - widely adopted in LLM fine-tuning for stable convergence. We save model checkpoints every 1k steps and monitor loss curves to ensure both Pilot and Copilot remain well-behaved.
• Consistent Performance under Different Learning Rates (LRs) In our hyperparameter tuning, we have tested different learning rates and observed consistent performance below: |LRs of Pilot (Qwen2.5-7B)|LRs of Copilot (T-Copilot-3B)|PIQA|HellaS.|SIQA| |-|-|-|-|-| |$3e^{-4}$|$1e^{-4}$|92.2|95.6|83.9| |$3e^{-4}$|$3e^{-4}$|92.2|95.8|84.2| |$3e^{-4}$|$5e^{-4}$|92.5|95.3|84.3|

---

[Limitation] Plug-and-Play Flexibility.

We thank the reviewer for providing the detailed suggestions. Regarding the plug-and-play designs, we would like to clarify that:

• Transferability of Copilot: In Appendix G.3, our transferability experiments show that a Copilot trained with one Pilot can pair with other Pilots of the same type, achieving equal error-correction gains. This demonstrates that Copilot is not tied to a single Pilot and can be used plug-and-play with similar models without retraining.
• Lightweight Integration with PEFT: In experiments, we show that T-Copilot is effective with both full fine-tuning and PEFT (e.g., LoRA) on LLaMA-3 and Qwen2.5 Pilots. PEFT updates only a small subset of parameters in Pilot and Copilot, substantially reducing training compute.
• Additional Adoption: We will explore more efficient, plug‑and‑play designs of T‑Copilot (e.g., eliminating the need for any Pilot model retraining) to broaden adoption in our future work.

Dear Reviewer 2chS,

As the discussion period is nearing its end, would you mind reading our response and letting us know if your concerns are addressed? Thank you so much for your time and consideration!

Warm Regards,

The Authors of Transformer Copilot

Dear Reviewer 2chS,

Thank you for raising score! We are pleased that our rebuttal has addressed your concerns. We will incorporate the results and analyses from the rebuttal to further strengthen our paper.

Warm Regards,

The Authors

Thank you for your constructive and thoughtful suggestions and feedback! We will address each of your concerns below.

---

[W.1] OOD Generalization of Mistake Log.

Thank you for the insightful question on the effectiveness of T-Copilot on the OOD data. In response:

1. T-Copilot's Design on Generality.

• Mistake Log Encodes Data-Agnostic Error Patterns: The Mistake Log systematically records model-internal signals (e.g., hidden states representations) besides dataset-specific idiosyncrasies, enabling the Copilot to learn the Pilot’s inherent error patterns.
• Initialization from pre-trained backbones enhances generality:
  As noted in Section 3.1, the Copilot is initialized from the same pre-trained backbone as the Pilot, leveraging broad language knowledge and helping avoid overfitting to any single fine-tuning dataset.

2. New Cross-Benchmark (OOD setting) Evaluation.
For a better demonstration, we fine-tune T-Copilot on arithmetic reasoning data (Math10k) and evaluate it on commonsense reasoning tasks (PIQA, HellaSwag, BoolQ, and OBQA) to assess T-Copilot's performance in OOD settings.

• From the results, although the Pilot model (LLaMA-3.1-8B) achieves less accuracy in the OOD setting compared to in-domain fine-tuning (Table 2 in paper), integrating T-Copilot-3B still yields consistent performance gains to the Pilot model.
• This suggests that T-Copilot remains effective in OOD settings due to its strong error-correction capability.

3. Limitations and Future Work.
We also acknowledge, as noted in Appendix A (Lines 620–621), that T-Copilot may face limitations under large domain shifts. We have explicitly identified improving its generalizability to such settings as an important direction for future work.

---

[W.2] Temporal Discrepancy in Mistake Log Sampling.

Thank you for mentioning this thoughtful point regarding the potential temporal drift in error relevance as the Pilot model evolves during fine-tuning.

The main objective of the Mistake Log, which collects and maintains the Pilot’s evolving hidden states, is to:

• enrich the diversity of error patterns;
• enable the Copilot to learn from and handle a broad range of error signals during inference.

During the early design of the Mistake Log, we identified two key challenges:

• First, as noted by the reviewer, randomly sampling from the full training history may introduce temporal discrepancies, leading to potential instability in Copilot training.
• Second, in practice, it might be infeasible to store the entire Mistake Log on GPU due to memory constraints.

To address these issues, we correspondingly introduce the Mistake Log Buffer (Lines 892–899), which retains only the most recent set of training rounds (e.g., 128). This design ensures that:

• The Copilot is trained on relatively recent Pilot states, striking a balance between reducing temporal discrepancies and avoiding overfitting to the latest state;
• The buffer can be efficiently stored and accessed on GPU, improving the method’s practicality.

We hope our explanations of the Mistake Log design address the reviewer’s question. We will also incorporate these rationales into our paper to improve clarity.

Dear Reviewer,

Thank you very much for your positive score of our submission. We will incorporate the clarifications and your thoughtful feedback to further strengthen our paper.

Warm regards,

The Authors of Transformer Copilot