In-Context Deep Learning via Transformer Models

Thanks for your detailed review. We have revised our draft (changes marked in BLUE) in this anonymous Dropbox folder.

C1: Reconstruction of the Manuscript.

Response

Thank you for the suggestions. We agree with this. In response, we have made the following modifications in the revised version:

• We have moved the definitions of transformers and attention into the introduction while keeping the full standard definition in Appendix D due to space constraints.

• The results on softmax attention and the experimental results have been incorporated into the main body—softmax results in Section 3.4 and experimental results in Section 4.

• To improve clarity, we simplify Problem 2 and clearly state that it describes an MLP with $N$ layers.

• The proof sketches in Section 3.3 have been moved to Appendix D, just before the detailed proof of the lemmas.

• The complex constants in Lemma 6 have been shifted to Appendix D to improve readability.

• While we acknowledge that Definition 1 may appear verbose, it establishes the notation used throughout the paper, particularly in the proofs, ensuring precision and consistency. Therefore, we have retained it as is.

Additionally, we appreciate the suggestion to include a visualization of the backpropagation process. We have included it in Appendix A.2 of the revised version. We also redisplay it anonymous figure.

C2: Difference with Standard Transformer

Response

Thank you for your suggestions. We have added this to the last point of limitations (line 749-754) and restate it here:

There are two minor differences between the transformer used in the theoretical analysis and a standard transformer: (i) The transformer used in the theoretical analysis incorporates an element-wise multiplication layer, a specialized variant of self-attention that retains only the diagonal score and allows efficient implementation. (ii) It does not alternate self-attention and MLP layers. We emphasize that this also qualifies as a standard transformer because we view either an attention or an MLP layer as equivalent to an attention plus MLP layer due to the residual connections.

C3 Example of Explicit Construction of ReLU

Response

Thank you for your insightful comments. Here, we provide clarifications on why the constants can be explicitly instantiated and illustrate this using a simple example.

The key reason is that the function approximated by the sum of ReLUs is relatively simple in our context, such as the Sigmoid activation function. For such simple functions, it is straightforward to derive an explicit construction.

Here, we take the Sigmoid activation function as an example and propose one explicit construction method. Let $r(z)$ denote the Sigmoid function.

1. Segment the input domain. For example, divide the domain $[-10, 10]$ smaller intervals such as $[-10, -9], [-9, -8], \dots, [9, 10]$.

2. Approximate each segment locally using a linear function via linear interpolation. For instance, in the domain $[9,10]$, approximate $r(z)$ using a linear function $a_1 z + c_1$, where $a_1$ and $c_1$ are calculated as follows: (i) $a_1 = (r(10)-r(9)) / (10-9)$. (ii) $c_1 = r(9) - a_1 * 9$.

3. Approximate linear function $a_1 z + c_1 (z \in [9,10])$ using a sum of ReLU terms. This step involves two substeps, which are straightforward to implement: (i) Approximate the indicator function for $z \in [9,10]$ using a sum of ReLU terms. (ii) Approximate the constant $c_1$ using the sum of ReLU. This is because bias terms are not included in the sum of ReLU terms in Definition 4. The bias term $c_1$​ must be approximated using an additional sum of ReLU terms.

4. Combine all the sum of ReLU approximators across all segments. Finally, integrate the approximations for all segments to construct the complete approximation.

5. Estimation of the parameters in Definition 4. $\epsilon_{approx} = 0.625$, $R=10$, $H=80$, and $C=25$.

Furthermore, to achieve higher precision in the approximation, it is sufficient to use finer segmentations.

Q1: Difference of Dimensions

Response

Thank you for your question. We apologize for any confusion caused. Here are some clarifications:

• In Section 2.(i), the value of $D$ is derived from prior research on in-context learning (ICL) for functions, indicated as $D=\Theta(d)$. This basis does not pertain to our focus on models, such as an $N$-layer neural network. We have removed these parts in the revised version to make things easier to read.

• Regarding the zero values in the rows of q_i, these primarily consist of intermediate terms such as $\bar{p}_i(j), \bar{r}^{‘}_i(j), \bar{s}_t(j)$, which are essential for calculating partial derivatives and contain gradient information. These rows are initially set to zero.

Thanks for your detailed review. We have revised our draft and addressed all concerns. The revised version (changes marked in BLUE) is available in this anonymous Dropbox folder.

Q1: Do trained transformers actually achieve the construction?

Response:

Thank you for the question. The trained transformers do not always achieve the construction. We apologize for any confusion caused, and acknowledge the limitation you mentioned. However, this limitation does not affect the primary contributions of our work. Here are some clarifications.

The main contribution of this paper is to provide an explicit construction demonstrating the existence of a transformer capable of simulating gradient descent (GD) for $N$-layer neural networks via in-context learning (ICL). The experimental design empirically validates that the trained transformer can indeed simulate GD steps for $N$-layer neural networks, supporting our theoretical results on the existence of such a transformer. Although there is a discrepancy between the theoretically constructed transformer and the empirically used one, this difference does not weaken the central point of our work, which is establishing the theoretical existence of such a transformer.

We acknowledge the limitation you mentioned and have incorporated it into the revised version: line 744-748.

C1: I think the paper suffers a lot from reduced clarity. It seems that most of the content of the paper is in the appendix, including empirical comparisons as well as discussion, which to me makes for a less clear read. I would suggest the authors rewrite the main body of text to include more intuitions and takeaways (e.g. when some assumptions may hold in practice rather than just in theory).

Response:

Thank you for your thoughtful suggestions. We agree that the main text can be better organized to improve readability and self-containment. In response, we have made the following modifications in the revised version:

• The results on softmax attention and the experimental results have been incorporated into the main body—softmax results in Section 3.4 and experimental results in Section 4.

• The proof sketches have been moved to Appendix D, just before the detailed proof of the lemmas.

• We include a remark concerning the practicality of our assumptions (Remark 4: line 1577-1579). Our assumptions remain modest. For example, we require that the loss function $l$, the activation function $r$, and its derivative be $C^4$-smooth. This condition is met by numerous network architectures, including those using the sigmoid activation function $r$ and the squared loss function $l$.

We hope our responses address your concerns and look forward to further feedback.

Thank you for your review! We greatly appreciate your attention to detail and recognition of our theoretical and experimental contributions! Your constructive comments and encouraging words are also highly appreciated!

Q1: How would the bound vary if the transformer's size is allowed to increase beyond (2N+4)L? Can it be made polynomial, by making the transformer size grow exponential?

Response:

Thank you for your question. The bound does not change if the transformer's size exceeds $(2N+4)L$, and the error cannot be reduced to a polynomial level. Here are some clarifications:

• Each $(2N+4)$-layer transformer simulates one gradient descent (GD) step, and we stack such transformers $L$ times.

• In non-convex optimization, each step’s trajectory depends on the previous steps. This dependency causes the exponential accumulation of the error.

Q2: Can some form of approximate self-attention or MLP layers be used in their place? How would the size of the constructed model vary then?

Response:

Thank you for your insightful question. The answer is No. Here are some clarifications.

• Explicit construction of transformer: We aim to provide an explicit transformer construction that simulates $L$ gradient descent steps. Although self-attention or MLP layers can approximate Hadamard products through their universal approximation capabilities, such approximations would result in non-explicit constructions.
• Specialized variant of self-attention: EWML can be viewed as a specialized variant of self-attention that retains only the diagonal scores.

We hope our responses address your concerns and look forward to further feedback.