SPB Git

Repository Public Research preview

models_and_objectives.tex 16.5 KB 257 lines latex

1% ============================================================================2% Models and objectives3% ============================================================================4\providecommand{\ph}[1]{\texttt{[[#1]]}}5\providecommand{\evid}[1]{\texttt{\detokenize{#1}}}67\section{Models and Objectives}8\label{sec:models}910\subsection{Problem formulation}1112Let $y \in \mathcal{Y}$ be an original passage in placeholder form and $m$ its metadata (discipline, era of13publication, section role, source). A rewriting process $R$ maps the original to a rewrite14\begin{equation}15  x = R(y, v, \xi), \qquad v \in \mathcal{V} = \{\mathrm{A}, \mathrm{B}, \mathrm{C}, \mathrm{D}, \mathrm{E}\},16  \label{eq:rewrite}17\end{equation}18where $v$ is the rewrite instruction and $\xi$ the sampling randomness of the LLM. A reconstruction model is an19estimator $\hat y = f_\theta(x, m)$ trained on validated pairs $\{(x_i, y_i, m_i)\}$. At inference $f_\theta$ sees20only the rewrite and, optionally, metadata supplied by the caller; it never sees $y$, and in particular the21originals of held-out items are never used. The variant $v$ is not an input: the model must infer the kind of22rewrite from $x$.2324\paragraph{Non-identifiability.} The rewriting process is many-to-one: distinct originals can produce the same or25nearly the same rewrite, and the rewrite instruction deliberately discards surface form. Exact recovery of $y$ from26$x$ is therefore not identifiable in general, and the target of learning is the conditional distribution27$p(y \mid x, m)$ rather than a function. A simple bound makes the consequence concrete. Let $d$ be a metric on28passages (for example, edit distance) and suppose two originals $y \neq y'$ are mapped, with equal probability, to29the same rewrite $x$. For any deterministic estimator, the triangle inequality gives30$d(y, y') \leq d(f_\theta(x,m), y) + d(f_\theta(x,m), y')$, so the expected error on these two items is at least31$\tfrac12 d(y, y')$. The distances between originals whose rewrites are near-identical thus give a data-only lower32bound on the error of \emph{any} deterministic reconstruction, which we estimate in Section~\ref{sec:design}. Two33consequences guide the evaluation. Lower lexical overlap with a particular original is not, by itself, evidence of34failure to recover style. And because content is shared across the several rewrites of one original while their35surface form varies, the paired-view identifiability results of \citet{vonkugelgen2021selfsupervised} suggest that36content is recoverable while the original style is recoverable only to the extent that the rewrite retains traces37of it. This motivates separating a content representation from a style representation in M3.3839\subsection{Shared backbone and tokenizer}4041All HTR models share a pre-layer-normalisation encoder--decoder Transformer \citep{vaswani2017attention,42xiong2020layernorm} implemented in PyTorch.\footnote{\evid{src/htr/models/transformer.py},43\evid{src/htr/models/config.py}.} Token embeddings are scaled by $\sqrt{d}$ and tied with the output projection.44Self-attention in encoder and decoder uses rotary position embeddings \citep{su2021roformer}; cross-attention does45not. Each encoder layer computes46\begin{equation}47  h \leftarrow h + \mathrm{Attn}(\mathrm{LN}(h)), \qquad h \leftarrow h + \mathrm{FFN}(\mathrm{LN}(h)),48\end{equation}49and each decoder layer inserts a cross-attention sub-layer over the encoder memory between causal self-attention50and the feed-forward sub-layer. The feed-forward block is SwiGLU,51$\mathrm{FFN}(h) = W_2\big(\mathrm{SiLU}(W_1 h) \odot W_3 h\big)$, without biases. Final layer normalisations are52applied to the encoder output $H \in \mathbb{R}^{L \times d}$ and to the decoder states. The tokenizer is a53SentencePiece unigram model \citep{kudo2018sentencepiece} trained from scratch on training-split text only, with54atomic tokens for placeholders and control symbols; this avoids test leakage through the vocabulary and makes the55tokenizer part of the release.5657\subsection{Model family}5859\paragraph{M1 (plain).} The backbone, with the copy mechanism described below.6061\paragraph{M2 (metadata-conditioned).} Learned embeddings of discipline, era, section role and source are summed62and projected into a metadata vector63\begin{equation}64  \mu = W_\mu\big(e_{\mathrm{disc}}(m) + e_{\mathrm{era}}(m) + e_{\mathrm{role}}(m) + e_{\mathrm{src}}(m)\big) + b_\mu,65\end{equation}66which is added to every encoder and decoder input position and prepended to the encoder memory as one extra67memory token. As a control, M1 can instead receive the same metadata as text tokens prefixed to the source68(\texttt{<disc:$\cdot$> <era:$\cdot$> <role:$\cdot$> <src:$\cdot$>}).6970\paragraph{M3 (dual-latent).} M3 keeps the metadata vector $\mu$ added to the encoder and decoder inputs, replaces71M2's single metadata memory token by $T_s$ style memory tokens, and adds a content vector, a style branch and an72adversary. With $\bar h$ and $\check h$ the masked mean and max of the encoder output over non-padding positions,73\begin{align}74  c &= W_c \bar h + b_c && \text{(content vector)} \label{eq:content}\\75  s &= \mathrm{MLP}_s\big([\bar h; \check h] + [\mu; \mu]\big) \in \mathbb{R}^K, \quad \pi = \mathrm{softmax}(s) && \text{(style branch)}\\76  S &= \textstyle\sum_{k=1}^{K} \pi_k P_k \in \mathbb{R}^{T_s \times d}, \quad \sigma = W_\sigma \big(\tfrac{1}{T_s}\textstyle\sum_{t} S_t\big) + b_\sigma && \text{(style memory and style vector)}\\77  a &= \mathrm{MLP}_a\big(\mathrm{GRL}_\lambda(c)\big) \in \mathbb{R}^K && \text{(adversary)}78\end{align}79Here $P_k \in \mathbb{R}^{T_s \times d}$ are learned style prototypes, $T_s = 4$ style tokens, $\mathrm{MLP}_s$ is a80two-layer GELU network ($2d \to d \to K$), and $\mathrm{MLP}_a$ a two-layer network ($d \to d/2 \to K$). The81gradient-reversal layer $\mathrm{GRL}_\lambda$ is the identity in the forward pass and multiplies the gradient by82$-\lambda$ in the backward pass, with $\lambda = 0.1$ \citep{ganin2014unsupervised}. The decoder cross-attends to the83memory $[S; H]$, so the predicted style distribution enters decoding through $T_s$ prefix tokens. The $K$ style84classes are the (discipline, era) buckets with at least 20 training pairs, plus a residual ``other'' class; they are85derived from training metadata only. The style branch is supervised only during training, and at inference the86style distribution is predicted from the rewrite (and the metadata, when supplied). The content vector is87encouraged to be invariant across a passage's original and rewrite, the style branch to predict the style class,88and the adversary and an orthogonality penalty to keep style information out of $c$; whether this yields89disentanglement is tested by ablations and a post-hoc probe, not assumed \citep{elazar2018adversarial}.9091\paragraph{Copy mechanism.} Because most of the original is recoverable by copying from the rewrite, all models92use a pointer-style copy gate over source positions \citep{see2017pointer,gu2016copying}, which can be switched93off for ablation. With decoder state $h_t$ at step $t$, output projection $W_o$ (tied with the embeddings), and94$\tilde\alpha_{t,i}$ the head-averaged cross-attention weight of the last decoder layer on source position $i$,95restricted to source tokens (memory prefix tokens excluded) and renormalised to sum to one, the output96distribution is97\begin{equation}98  p(w \mid y_{<t}, x, m) = (1 - g_t)\,\mathrm{softmax}(W_o h_t)_w + g_t \sum_{i:\,x_i = w} \tilde\alpha_{t,i},99  \qquad g_t = \mathrm{sigmoid}(w_g^\top h_t + b_g).100  \label{eq:copy}101\end{equation}102The model returns $\log\big(p(\cdot) + 10^{-9}\big)$.103104\begin{figure}[t]105\centering106\fbox{\parbox{0.9\textwidth}{\centering\vspace{2em}\pending{architecture figure: rewrite $x$ (+ metadata $m$)107$\to$ encoder $E_c$ $\to$ content vector $c$ and memory $H$; style branch $E_s \to \pi \to$ style memory $S$;108gradient-reversal adversary on $c$; decoder cross-attending to $[S;H]$ with copy gate; training-only encoding109of the original $y$ for $c'$}\vspace{2em}}}110\caption{Schematic of the dual-latent HTR model (M3). M1 omits the metadata, style and adversarial components; M2111keeps the metadata vector but omits the style branch, adversary and content losses. The figure is generated by a112script from the model definition.}113\label{fig:arch}114\end{figure}115116\subsection{Training objective}117118Let $y_{1:T}$ be the target tokens, $p_{b,t} = p(\cdot \mid y_{<t}, x_b, m_b)$ the copy-augmented output119distribution of item $b$ in a batch of size $B$, $N$ the number of non-padding target tokens in the batch, and $V$120the vocabulary size. The objective implemented in the code is\footnote{\evid{src/htr/models/losses.py}.}121\begin{equation}122  \mathcal{L} = \lambda_{\mathrm{ce}} \mathcal{L}_{\mathrm{token}} + \lambda_{\mathrm{fact}} \mathcal{L}_{\mathrm{fact}}123  + \mathbb{1}[\text{M3}]\big(\lambda_{\mathrm{style}} \mathcal{L}_{\mathrm{style}} + \lambda_{\mathrm{adv}} \mathcal{L}_{\mathrm{adv}}124  + \lambda_{\mathrm{con}} \mathcal{L}_{\mathrm{con}} + \lambda_{\mathrm{sem}} \mathcal{L}_{\mathrm{sem}} + \lambda_{\mathrm{orth}} \mathcal{L}_{\mathrm{orth}}\big),125  \label{eq:objective}126\end{equation}127with the following terms.128129\paragraph{Token loss.} Label-smoothed negative log-likelihood, averaged over non-padding target tokens:130\begin{equation}131  \mathcal{L}_{\mathrm{token}} = \frac{1}{N} \sum_{b,t} \Big[(1-\varepsilon)\big(-\log p_{b,t}(y_{b,t})\big)132  + \varepsilon \cdot \frac{1}{V}\sum_{w=1}^{V}\big(-\log p_{b,t}(w)\big)\Big], \qquad \varepsilon = 0.1 .133\end{equation}134Smoothing is switched off ($\varepsilon = 0$) in the one-batch overfit test and in the validation loss used for135early stopping.136137\paragraph{Factual loss.} The unsmoothed negative log-likelihood restricted to factual target pieces. With138$f_{b,t} = 1$ if target piece $t$ contains a placeholder delimiter or a digit and $0$ otherwise,139\begin{equation}140  \mathcal{L}_{\mathrm{fact}} = \frac{\sum_{b,t} f_{b,t}\big(-\log p_{b,t}(y_{b,t})\big)}{\max\big(1, \sum_{b,t} f_{b,t}\big)} .141\end{equation}142It up-weights the spans that carry citations, equations, links and numbers; $\lambda_{\mathrm{fact}} = 0$ recovers143the plain token loss.144145\paragraph{Style loss.} Cross-entropy of the style branch against the training-only style class $k_b$:146$\mathcal{L}_{\mathrm{style}} = \frac{1}{B}\sum_b -\log \mathrm{softmax}(s_b)_{k_b}$.147148\paragraph{Adversarial loss.} Cross-entropy of the adversary that predicts the style class from the content149vector through the gradient-reversal layer:150$\mathcal{L}_{\mathrm{adv}} = \frac{1}{B}\sum_b -\log \mathrm{softmax}(a_b)_{k_b}$. Minimising $\mathcal{L}$151trains the adversary to predict $k_b$, while the reversed gradient pushes the content encoder to make $c$152uninformative about style.153154\paragraph{Contrastive and semantic losses.} Let $c'_b$ be the content vector of the \emph{original} $y_b$,155computed by the same encoder with the same metadata under stop-gradient ($\mathrm{sg}$), and let $\hat c$ denote156L2-normalised vectors. With the similarity matrix $Z_{ij} = \hat c_i^\top \mathrm{sg}(\hat c'_j)/\tau$, $\tau =1570.07$, and the other items of the batch as negatives,158\begin{equation}159  \mathcal{L}_{\mathrm{con}} = \tfrac12\big[\mathrm{CE}(Z, I) + \mathrm{CE}(Z^\top, I)\big]160  = -\frac{1}{2B}\sum_{b}\Big[\log\frac{e^{Z_{bb}}}{\sum_j e^{Z_{bj}}} + \log\frac{e^{Z_{bb}}}{\sum_j e^{Z_{jb}}}\Big],161\end{equation}162an InfoNCE objective \citep{oord2018cpc,chen2020simclr}, and163\begin{equation}164  \mathcal{L}_{\mathrm{sem}} = 1 - \frac{1}{B}\sum_b \cos\big(c_b, \mathrm{sg}(c'_b)\big),165\end{equation}166a cheap internal surrogate for an external embedding similarity. Neither term is a semantic metric of the output167text; they act on the encoder's content representation.168169\paragraph{Orthogonality.} With $\tilde C, \tilde\Sigma \in \mathbb{R}^{B \times d}$ the batch-centred content vectors170$c_b$ and style vectors $\sigma_b$,171\begin{equation}172  \mathcal{L}_{\mathrm{orth}} = \frac{1}{d^2}\,\Big\| \tfrac{1}{\max(1, B-1)}\,\tilde C^\top \tilde\Sigma \Big\|_F^2 ,173\end{equation}174the mean squared cross-covariance between content and style vectors \citep{bousmalis2016domain}.175176\paragraph{Differentiability.} Every term is differentiable with respect to the model parameters. The factual mask177$f$ and the style labels $k$ are constants; the copy mixture is differentiable through the gate and the attention178weights; $c'$ receives no gradient; and logged accuracies are outside the computation graph. The original is179encoded only inside the loss, during training.180181\paragraph{Weights.} Table~\ref{tab:loss-weights} gives the weight presets. The M3 weights are initial values; any182tuning uses the validation split only, and the token-loss-only configuration of the same architecture is reported183as a control. The ablation presets each remove one component.184185\begin{table}[t]186\centering187\caption{Loss-weight presets (\texttt{scripts/train.py}). A dash means the term is inactive (weight 0).}188\label{tab:loss-weights}189\small190\begin{tabular}{@{}lccccccc@{}}191\toprule192Preset & $\lambda_{\mathrm{ce}}$ & $\lambda_{\mathrm{fact}}$ & $\lambda_{\mathrm{style}}$ & $\lambda_{\mathrm{adv}}$ & $\lambda_{\mathrm{con}}$ & $\lambda_{\mathrm{sem}}$ & $\lambda_{\mathrm{orth}}$ \\193\midrule194M1, M2 & 1 & -- & -- & -- & -- & -- & -- \\195M3 & 1 & 0.5 & 0.3 & 0.3 & 0.2 & 0.1 & 0.1 \\196M3, token loss only & 1 & -- & -- & -- & -- & -- & -- \\197M3, no style & 1 & 0.5 & -- & -- & 0.2 & 0.1 & -- \\198M3, no adversary & 1 & 0.5 & 0.3 & -- & 0.2 & 0.1 & 0.1 \\199M3, no factual & 1 & -- & 0.3 & 0.3 & 0.2 & 0.1 & 0.1 \\200M3, no contrastive/semantic & 1 & 0.5 & 0.3 & 0.3 & -- & -- & 0.1 \\201\bottomrule202\end{tabular}203\end{table}204205\subsection{Size design points}206207Table~\ref{tab:sizes} lists the design points of the model family. The parameter counts are those of an M1 model208at initialisation with a 32\,000-token vocabulary; M2 and M3 add a small number of parameters (metadata embeddings,209style branch, prototypes and adversary). They are design points, not promises: the counts of trained checkpoints,210which depend on the trained tokenizer's vocabulary, are logged in each run manifest and are \pending{measured211parameter counts of trained checkpoints}. HTR-Base is the primary release target; HTR-Large is conditional on212measured feasibility.213214\begin{table}[t]215\centering216\caption{Size design points (\texttt{htr.models.config.SIZES}). Parameters: M1 at initialisation, vocabulary21732\,000, tied embeddings.}218\label{tab:sizes}219\small220\begin{tabular}{@{}lrrrrr@{}}221\toprule222Design point & $d_{\mathrm{model}}$ & Heads & $d_{\mathrm{ff}}$ & Enc./dec. layers & Parameters \\223\midrule224HTR-Tiny & 384 & 6 & 1\,536 & 4 / 4 & $\approx$ 33\,M \\225HTR-Small & 768 & 12 & 3\,072 & 5 / 5 & $\approx$ 130\,M \\226HTR-Base & 1\,024 & 16 & 4\,096 & 8 / 8 & $\approx$ 335\,M \\227HTR-Large & 1\,536 & 24 & 6\,144 & 12 / 12 & $\approx$ 1.07\,B \\228\bottomrule229\end{tabular}230\end{table}231232\subsection{Baselines}233234\paragraph{B0 (identity).} The output is the rewrite itself. Because the rewrites preserve content by construction,235B0 is a strong baseline on every overlap and factual metric, and any claimed gain is stated relative to it.236237\paragraph{B1 (generic paraphrase).} A deterministic rule-based paraphraser (the project's mock rewriter, variant A)238applied to the rewrite. It tests whether a gain could come from \emph{any} further rewriting rather than from a239learned inverse. It is weak by design because it does not learn; a minimally trained rewrite baseline is240\pending{trained generic-rewrite baseline}.241242\paragraph{B2 (pretrained encoder--decoder).} Fine-tuning a pretrained T5-family model \citep{raffel2019t5} on the243same pairs and splits, with metadata given as a text prefix and outputs passed through the same placeholder244protocol.\footnote{\evid{src/htr/training/hf_finetune.py}; BART \citep{lewis2019bart} is an alternative.} B2 is a245secondary comparator: its pretraining data may contain the originals, so it is subject to the contamination audit of246Section~\ref{sec:design}. The base model and its licence are \pending{B2 base model}.247248\subsection{Decoding and inference}249250Decoding is deterministic: greedy, or beam search with a length penalty, with key--value caching and the copy251mechanism active; the beam width, if any, is chosen on validation and fixed in the pre-registration. The released252inference interface protects citations, mathematics and URLs in the input, tokenises it (with the metadata prefix253when the model was trained with one), decodes, restores the placeholders, and attaches preservation warnings for254changed numbers, placeholders, negation, hedging or length.\footnote{\evid{src/htr/inference/pipeline.py}.} Every255result carries the notice that it is ``a model's estimate of pre-2018-style scientific prose given an LLM rewrite;256it is not the original text and is not evidence of human authorship''.257
Files main
Menu

Theme