Skip to main content

MATH 345: Linear Algebra and Optimization

Section 5.7 Fixed-hidden-layer training

Let \(N_{\mathbf{c}}(t)=c_0+c_1\sigma(t)+c_2\sigma(t-1)+c_3\sigma(t+1)\text{,}\) where \(\sigma(t)=\tanh(t)\text{.}\) Only \(c_0,c_1,c_2,c_3\) are trainable.
For example, take data
\begin{equation*} t=(-2,-1,0,1,2),\qquad y=(4,1,0,1,4). \end{equation*}
Given data \((t_i,y_i)\text{,}\) define
\begin{equation*} A=\begin{bmatrix} 1&\sigma(t_1)&\sigma(t_1-1)&\sigma(t_1+1)\\ 1&\sigma(t_2)&\sigma(t_2-1)&\sigma(t_2+1)\\ \vdots&\vdots&\vdots&\vdots\\ 1&\sigma(t_m)&\sigma(t_m-1)&\sigma(t_m+1) \end{bmatrix}. \end{equation*}
Training the final layer is the least-squares problem \(A\mathbf{c}\approx\mathbf{y}\text{.}\) If the hidden-layer weights were also trainable, the loss would become nonlinear in the parameters, and gradient descent plus the chain rule would be needed.

Activity 5.7.1. Fixed features, trained coefficients.

Let \(\sigma(t)=\tanh(t)\text{,}\) and let
\begin{equation*} N_{\mathbf{c}}(t)=c_0+c_1\sigma(t)+c_2\sigma(t-1)+c_3\sigma(t+1). \end{equation*}
For data inputs \(t_1,\ldots,t_m\text{,}\) define
\begin{equation*} A= \begin{bmatrix} 1&\sigma(t_1)&\sigma(t_1-1)&\sigma(t_1+1)\\ 1&\sigma(t_2)&\sigma(t_2-1)&\sigma(t_2+1)\\ \vdots&\vdots&\vdots&\vdots\\ 1&\sigma(t_m)&\sigma(t_m-1)&\sigma(t_m+1) \end{bmatrix}. \end{equation*}
  1. Which entries of \(A\) are fixed once the data inputs are known?
  2. Which vector is trained?
  3. Why is \(A\mathbf{c}\approx \mathbf{y}\) a least-squares problem?
  4. Why would the problem become nonlinear if the shifts inside \(\sigma\) were trainable?
Tags. [U5-LO6, U5-LO7 | C+M+R | Core]
Solution.
Once the data inputs \(t_1,\ldots,t_m\) are known, every entry of \(A\) is fixed. The trained vector is
\begin{equation*} \mathbf{c}= \begin{bmatrix} c_0\\ c_1\\ c_2\\ c_3 \end{bmatrix}. \end{equation*}
The prediction vector has the form \(A\mathbf{c}\text{,}\) so choosing \(\mathbf{c}\) to make \(A\mathbf{c}\) close to \(\mathbf{y}\) is a least-squares problem. If the shifts inside \(\sigma\) were trainable, then the entries of \(A\) would depend on trainable parameters. The model would no longer be linear in all of the parameters being trained.

Note 5.7.1.

The model is nonlinear as a function of the input \(t\text{,}\) but linear as a function of the trained coefficient vector \(\mathbf{c}\text{.}\) This is why least squares applies to the fixed-hidden-layer problem.
For the five data points above, this produces a \(5\times4\) design matrix. The lab asks you to build that matrix, solve the least-squares problem for \(\mathbf{c}\text{,}\) and compare the fitted curve with the data.

Note 5.7.2. Machine-learning example: next-token prediction as optimization.

A language model receives a sequence of tokens and produces scores for possible next tokens. After converting scores into probabilities, training rewards the model for assigning high probability to the actual next token.
A final hidden vector \(\mathbf{h}\) can be converted into scores by an affine rule
\begin{equation*} \boldsymbol{\ell}=W\mathbf{h}+\mathbf{b}. \end{equation*}
These scores are often called logits. A later nonlinear step converts scores into probabilities, and the loss measures how well the model assigns probability to the correct next token.
For a sequence \(t_1,t_2,\ldots,t_L\text{,}\) a simplified next-token loss is
\begin{equation*} L(\theta) = -\sum_{i=1}^{L-1} \log p_\theta(t_{i+1}\mid t_1,\ldots,t_i). \end{equation*}
The details of a real language model are complicated, but the optimization idea is familiar: \(\theta_{k+1}=\theta_k-\alpha\nabla L(\theta_k)\text{.}\) The parameters \(\theta\) include many matrices and bias vectors. Training changes those parameters to reduce the loss. This is why gradients, chain rules, and matrix products are central in modern machine learning.

Activity 5.7.2. Reading a next-token loss as optimization.

A simplified language-model training objective can be written
\begin{equation*} L(\theta)=-\sum_{i=1}^{L-1}\log p_\theta(t_{i+1}\mid t_1,\ldots,t_i). \end{equation*}
A gradient descent step has the form
\begin{equation*} \theta_{k+1}=\theta_k-\alpha\nabla L(\theta_k). \end{equation*}
  1. What are the parameters?
  2. What quantity is being minimized?
  3. Which symbol is the learning rate?
  4. What does the gradient point toward?
  5. Why is this an optimization problem, even though we are not studying full model training?
Tags. [U5-LO2, U5-LO3 | C+M+R | Core]
Solution.
The parameters are collected in \(\theta\text{.}\) The objective \(L(\theta)\) is the loss. The learning rate is \(\alpha\text{.}\) The gradient \(\nabla L(\theta_k)\) gives the local direction of steepest increase, so the negative gradient is used for descent. This is an optimization problem because training means adjusting parameters to reduce a loss.

Activity 5.7.3. Training Only the Last Layer.

Suppose a fixed hidden representation is \(\mathbf{h}\in\mathbb R^d\text{,}\) and an output layer computes \(\ell=W\mathbf{h}\text{.}\) For one training example, suppose the loss gradient with respect to the logits is \(\mathbf{g}\in\mathbb R^m\text{.}\) Then the gradient with respect to \(W\) has the form
\begin{equation*} \nabla_W L = \mathbf{g}\mathbf{h}^T. \end{equation*}
What is the shape of \(\mathbf{g}\mathbf{h}^T\text{?}\) Why is this a rank-one matrix, unless \(\mathbf{g}=\mathbf{0}\) or \(\mathbf{h}=\mathbf{0}\text{?}\) What does the update \(W_{\mathrm{new}}=W-\alpha \mathbf{g}\mathbf{h}^T\) change?
Solution.
If \(\mathbf{g}\in\mathbb R^m\) and \(\mathbf{h}\in\mathbb R^d\text{,}\) then \(\mathbf{g}\mathbf{h}^T\) is \(m\times d\text{.}\) It is an outer product, so it has rank at most one. The update changes the rows of \(W\) in proportion to the output error and the active hidden coordinates.