\resizebox{\linewidth}{!}{%
\begin{minipage}{1.2\linewidth}
\foreach \pdf/\title [count=\i] in {
    {image/repeat-short2.pdf}/{(a) Repeat (KVzip)},
    {image/qa-short2.pdf}/{(b) QA-1},
    {image/qa-1-short2.pdf}/{(c) QA-2},
    {image/summary-short2.pdf}/{(d) Summarization},
    {image/self-short2.pdf}/{(e) Prefill ($\text{H}_2\text{O}$})} {%
    \begin{tikzpicture}
        % Load PDF heatmap
        \node[inner sep=0pt] (img\i) {\includegraphics[width=15cm,height=3cm,clip,trim=0 0 0 0]{\pdf}};

        % Title
        \node[above=-2mm of img\i.north, font=\bfseries] {\title};

        % Define the labels
        \begin{axis}[
            at={(img\i.south west)},
            anchor=south west,
            width=15cm,
            height=3cm,
            scale only axis,
            xmin=0, xmax=1, % Explicitly define to avoid warning
            xlabel={KV sequence (context)},
            xlabel shift=-0.15cm,
            label style={font=\footnotesize},
            colormap/viridis,
            colorbar,
            ymin=0.0,
            ymax=1.0,
            point meta min=0,
            point meta max=1,
            axis line style={draw=none},
            colorbar style={
            width=0.2cm,
            height=2.7cm,
            at={(1.0,0.5)},
            anchor=west,
            ytick={0,0.2,0.4,0.6,0.8,1},
            yticklabel style={
                /pgf/number format/.cd,
                fixed,
                fixed zerofill,
                precision=1
            },
            font=\footnotesize,
            },
            xtick=\empty,
            ytick={0.125,0.375,0.625,0.875},
            yticklabels={layer 24, layer 16, layer 8, layer 0},
            tick label style={font=\footnotesize},
            tick style={draw=none},
            yticklabel style={xshift=4pt,font=\footnotesize},
        ]
        % Empty plot, just to use colorbar feature
        \addplot [draw=none] coordinates {(0,0)};
        \end{axis}
    \end{tikzpicture}\par\vspace{0.4cm}
}
\end{minipage}
}