docs/method.tex

\section{Method}
\label{sec:method}

We introduce KVPath, a KV representation paradigm that connects
key and value states across network depth. Our central hypothesis is that
changes between neighboring layers' KV states can be represented in fewer
dimensions than the full states. KVPath exploits this structure by
representing a sequence of layer-wise KV states through an initial state
and compact descriptions of their changes.

\subsection{Path Representation}
\label{sec:kvpath_framework}

We partition the model's $L$ layers into groups of $m$ consecutive layers,
assuming $m$ divides $L$.
Within each group, the KV states of a token form a \emph{path} of length
$m$, with local layer indices $\ell=0,\ldots,m-1$. 
For each KV head, let $k^{(\ell)},v^{(\ell)}\in\mathbb{R}^{d}$ denote
the states before positional encoding. The core of KVPath is to express
each layer's KV state through the preceding state and a cross-layer change:
\begin{equation}
\begin{aligned}
    k^{(\ell)} &= k^{(\ell-1)}+\Delta k^{(\ell)},\\
    v^{(\ell)} &= v^{(\ell-1)}+\Delta v^{(\ell)},
    \qquad 1\leq\ell<m.
\end{aligned}
\label{eq:kvpath_general}
\end{equation}
KVPath represents these changes using fewer dimensions than the full KV
states, with a compact dimension $d_\delta<d$ per key or value head. Below, we describe
two ways to realize this compact representation, together with their
state construction and cache requirements.

\subsection{Path Construction}
\label{sec:path_constructions}

We realize the compact cross-layer changes in two ways, illustrated in
Figure~\ref{fig:kvpath}. Both begin with an \emph{alpha component} generated from the first
layer's hidden state and use a low-dimensional \emph{delta component}
to describe variation along the path. In KVPath-Concat, the alpha component occupies shared coordinates
of every KV state, while the delta component supplies the remaining,
layer-dependent coordinates. In KVPath-Residual, the alpha component
is the full initial KV state, while each subsequent layer's delta
component encodes an update that is expanded and added to the preceding
state. 

\begin{figure}[t]
    \centering
    \includegraphics[width=\linewidth]{figs/fig1.pdf}
    \caption{Two KVPath constructions within a four-layer path.
    (a) KVPath-Concat: each layer concatenates an identical alpha
    component with its own delta component. Cross-layer differences are
    confined to the coordinates of the delta component.
    (b) KVPath-Residual: the alpha component initializes the full K/V
    state. Subsequent layers project their inputs into low-dimensional
    K/V updates, expand them separately, and add them to the preceding state.}
    \label{fig:kvpath}
\end{figure}

\paragraph{KVPath-Concat: variation in fixed coordinates.}
As shown in Figure~\ref{fig:kvpath}(a), each KV state consists of a
shared alpha component $(k_\alpha,v_\alpha)$ with $d_\alpha$ dimensions
per key or value head and a layer-dependent delta component
$(k_\delta^{(\ell)},v_\delta^{(\ell)})$ with $d_\delta$ dimensions.
The first layer generates the alpha component through a path-specific projection
$W_\alpha$. Each layer, including layer $0$, generates its delta
component through its own projection $W_\delta^{(\ell)}$:
\begin{equation}
\begin{aligned}
    (k_\alpha,v_\alpha)&=W_\alpha h^{(0)},\\
    (k_\delta^{(\ell)},v_\delta^{(\ell)})
        &=W_\delta^{(\ell)}h^{(\ell)},
        \qquad 0\leq\ell<m.
\end{aligned}
\label{eq:concat_projections}
\end{equation}
Every layer reuses the alpha component without modification and concatenates it with
its own delta component to form the full state:
\begin{equation}
\begin{aligned}
    k^{(\ell)}&=[k_\alpha\,\Vert\,k_\delta^{(\ell)}],\\
    v^{(\ell)}&=[v_\alpha\,\Vert\,v_\delta^{(\ell)}],
\end{aligned}
\label{eq:kv_path}
\end{equation}
where $\Vert$ denotes concatenation and $d=d_\alpha+d_\delta$.
Since the alpha component is identical across the path, the cross-layer increments
in Eq.~\eqref{eq:kvpath_general} are
\begin{equation}
\begin{aligned}
    \Delta k^{(\ell)}
        &=[0\,\Vert\,k_\delta^{(\ell)}-k_\delta^{(\ell-1)}],\\
    \Delta v^{(\ell)}
        &=[0\,\Vert\,v_\delta^{(\ell)}-v_\delta^{(\ell-1)}],
        \qquad 1\leq\ell<m.
\end{aligned}
\label{eq:path_update}
\end{equation}
Each full-dimensional change is determined by a $d_\delta$-dimensional
difference, with zeros in the coordinates of the alpha component.

We cache the alpha component once per path and the delta component
for each of the $m$ layers. Counting
both K and V, the resulting storage is
\begin{equation}
    \mathrm{Cache}=2d_\alpha+2m d_\delta.
\label{eq:concat_cache}
\end{equation}

\paragraph{KVPath-Residual: cumulative low-dimensional updates.}
As shown in Figure~\ref{fig:kvpath}(b), the alpha component now denotes the entire
initial state $(k_\alpha,v_\alpha)$, with $d$ dimensions per key or
value head. Each subsequent layer produces a delta component
$(k_\delta^{(\ell)},v_\delta^{(\ell)})$, consisting of separate
$d_\delta$-dimensional K and V update coordinates. A path-specific projection
$W_\alpha$ initializes the state, while layer-specific projections
$W_\delta^{(\ell)}$ generate the compact updates:
\begin{equation}
\begin{aligned}
    (k^{(0)},v^{(0)})&=(k_\alpha,v_\alpha)=W_\alpha h^{(0)},\\
    (k_\delta^{(\ell)},v_\delta^{(\ell)})
        &=W_\delta^{(\ell)}h^{(\ell)},
        \qquad 1\leq\ell<m.
\end{aligned}
\label{eq:residual_anchor}
\end{equation}
The learned matrices $U_K^{(\ell)},U_V^{(\ell)}\in
\mathbb{R}^{d\times d_\delta}$ expand these coordinates separately
for K and V. Each layer adds the expanded updates to the preceding
full state:
\begin{equation}
\begin{aligned}
    k^{(\ell)}&=k^{(\ell-1)}+U_K^{(\ell)}k_\delta^{(\ell)},\\
    v^{(\ell)}&=v^{(\ell-1)}+U_V^{(\ell)}v_\delta^{(\ell)},
    \qquad 1\leq\ell<m.
\end{aligned}
\label{eq:residual_update}
\end{equation}
Accordingly, the cross-layer increments in Eq.~\eqref{eq:kvpath_general}
are
\begin{equation}
\begin{aligned}
    \Delta k^{(\ell)}&=U_K^{(\ell)}k_\delta^{(\ell)},\\
    \Delta v^{(\ell)}&=U_V^{(\ell)}v_\delta^{(\ell)}.
\end{aligned}
\label{eq:residual_delta}
\end{equation}
Each increment can affect all $d$ coordinates but lies in a subspace
of dimension at most $d_\delta$.

A compact representation stores the full alpha component once and the
delta components for the $m-1$ subsequent layers:
\begin{equation}
    \mathrm{Cache}=2d+2(m-1)d_\delta.
\label{eq:residual_cache}
\end{equation}

\subsection{Attention with KVPath}
\label{sec:kvpath_attention}

For KVPath-Concat, the query is split into an alpha component
$q_\alpha^{(\ell)}$ and a delta component $q_\delta^{(\ell)}$,
with dimensions $d_\alpha$ and $d_\delta$, respectively.
RoPE is applied only to the delta components of queries and keys,
while the alpha components remain unrotated:
\begin{equation}
\begin{aligned}
    (q_\alpha^{(\ell)},q_\delta^{(\ell)})
        &=W_Q^{(\ell)}h^{(\ell)},\\
    q^{(\ell)}
        &=[q_\alpha^{(\ell)}\,\Vert\,
          \operatorname{RoPE}(q_\delta^{(\ell)})],\\
    k^{(\ell)}
        &=[k_\alpha\,\Vert\,
          \operatorname{RoPE}(k_\delta^{(\ell)})],\\
    v^{(\ell)}
        &=[v_\alpha\,\Vert\,v_\delta^{(\ell)}].
\end{aligned}
\label{eq:kvpath_qkv}
\end{equation}
For KVPath-Residual, full-width RoPE is applied to the projected query
and reconstructed key before attention; the resulting representations
are denoted by $q^{(\ell)}$ and $k^{(\ell)}$.

For a given query head, $Q^{(\ell)}$ stacks its token-wise query
vectors $q^{(\ell)}$ as rows. The matrices $K^{(\ell)}$ and
$V^{(\ell)}$ similarly stack the keys and values of the KV head
assigned to that query head. Both constructions then use standard
causal attention:
\begin{equation}
    O^{(\ell)}=\operatorname{softmax}\!\left(
        \frac{Q^{(\ell)}(K^{(\ell)})^\top}{\sqrt{d}}+M
    \right)V^{(\ell)},
\label{eq:attention}
\end{equation}
where $M$ is the causal mask and softmax is applied row-wise.

KVPath constructs each path independently for each KV head and does
not assume a particular mapping between query heads and KV heads.
It therefore applies directly to multi-head attention (MHA), where
each query head has its own KV head; grouped-query attention (GQA),
where a group of query heads shares one KV head; and multi-query
attention (MQA), where all query heads share a single KV head. In
every case, KVPath changes how the states of each KV head are represented
across layers, while the standard assignment from query heads to KV
heads remains unchanged.

KVPath is also applicable to multi-head latent attention (MLA), because
the path can be constructed over MLA's compressed latent representation.
Our KVPath-MLA instantiation shares an alpha latent within each path and
uses layer-dependent components to retain local variation. We describe
this instantiation and report its results in
Section~\ref{sec:mla-compatibility}.