Compute‐Optimal Model Design

• Compute‐optimal model design is a rigorous framework that jointly optimizes architecture size, data allocation, and fine-tuning strategies under fixed compute budgets.
• It leverages empirical scaling laws and cost accounting to systematically select optimal configurations, balancing training compute with error minimization.
• The approach provides a clear recipe and phase-diagram analysis that guides the choice between full fine-tuning and efficient methods like LoRA.

Compute‐optimal model design encompasses a mathematically rigorous framework for selecting architecture size, data volume, and fine-tuning protocol that jointly minimizes error (e.g., final task loss) under a strict compute (FLOP) budget. In the context of repurposing pre-trained decoder-only LLMs into contrastive embedding models, Ziarko et al. provide the first explicit parameterization and enumeration of compute-optimal fine-tuning strategies, model scaling, and data allocation, grounded in empirical scaling laws, analytic cost accounting, and a discrete phase-diagram of fine-tuning algorithms (Ziarko et al., 2024).

1. Formal Optimization Problem

The compute-optimal embedding model design problem is cast as a constrained minimization:

$$\begin{aligned} &\text{minimize}_{N, D, M,\ldots}\quad L(N, D, M,\ldots) \ &\text{subject to}\hspace{2em} C(N, D, M,\ldots) \leq C_\text{budget} \end{aligned}$$

Decision variables:

• $N$: non-embedding model parameters (model size)
• $D$: number of tokens used for contrastive fine-tuning
• $M \in \{\text{full, freeze, bias, LoRA}\}$: fine-tuning method
• $s \in [0, 1]$: fraction of active transformer blocks (for block-freezing)
• $r \in \mathbb{N}$: rank in LoRA adaptation

Here, $L(N, D, M, ...)$ is the final (contrastive) loss and $C(N, D, M, ...)$ is the total incurred FLOP count. The goal is to select the configuration achieving the best embedding quality for a given training compute.

2. Compute-Cost and Scaling Law Formulation

Compute cost:

Let $N_F$ and $N_B$ be the numbers of parameters participating in forward and backward passes per token, respectively. The total training compute for $N$0 tokens is

$N$1

which reduces to $N$2 in full fine-tuning (all parameters active during forward and backward). For partial/freezing/LoRA, $N$3 and $N$4 are correspondingly reduced.

Empirical scaling laws:

The final embedding error (contrastive loss) typically follows a power law over both $N$5 and $N$6:

$N$7

where $N$8 are empirically fitted, and $N$9 is the irreducible loss floor. This law is central to compute-optimality: it quantifies diminishing returns to model and data scaling.

3. Search Algorithm for Optimal Configuration

Given a set of pretrained LMs indexed by $D$0, a grid of data budgets $D$1 (total tokens), and candidate fine-tuning schemes $D$2, Ziarko et al. introduce a brute-force search over the grid:

$r \in \mathbb{N}$0

This systematic grid search is justified by the low dimensionality (few model sizes and data points) and discrete fine-tuning choices. Moreover, block freezing fraction $D$3 and LoRA rank $D$4 are similarly swept for their respective methods.

4. Fine-Tuning Method Phase Diagram

Empirical results reveal a phase diagram in which the compute-optimal fine-tuning method varies with budget:

• Full fine-tuning dominates at limited compute ($D$5 FLOP); it utilizes all parameters for maximum flexibility when few updates are available.
• LoRA (Low-Rank Adaptation) with moderate-to-high rank $D$6 becomes optimal at larger budgets ($D$7 FLOP), leveraging parameter-efficient adaptation for continued gains as data quantity grows large.
• Freeze or bias-only methods are never optimal at any realistic budget, as their performance saturates rapidly with compute.

Optimal adaptation settings (block-freezing ratio $D$8, LoRA rank $D$9) depend smoothly on $M \in \{\text{full, freeze, bias, LoRA}\}$0. The phase boundary is empirically sharp, cross-validated on multiple embedding tasks.

5. Compute-Optimal Model Design Recipe

For practitioners, the following explicit recipe emerges:

• Step 1: Choose the largest pretrained LM $M \in \{\text{full, freeze, bias, LoRA}\}$1 feasible under your compute budget, subject to cost constraints for the planned fine-tuning method.
• Step 2: Allocate as much fine-tuning data $M \in \{\text{full, freeze, bias, LoRA}\}$2 as can fit with the selected $M \in \{\text{full, freeze, bias, LoRA}\}$3 and method $M \in \{\text{full, freeze, bias, LoRA}\}$4 while not exceeding $M \in \{\text{full, freeze, bias, LoRA}\}$5.
• Step 3: For $M \in \{\text{full, freeze, bias, LoRA}\}$6 below $M \in \{\text{full, freeze, bias, LoRA}\}$7–$M \in \{\text{full, freeze, bias, LoRA}\}$8 FLOP, prefer full fine-tuning; above this threshold, switch to LoRA with rank $M \in \{\text{full, freeze, bias, LoRA}\}$9 maximizing downstream performance within $s \in [0, 1]$0.
• Step 4: Never select parameter-freezing or bias-only unless explicitly compute-constrained beyond the points evaluated in (Ziarko et al., 2024) (as they underperform).
• Step 5: For block freezing in “full” fine-tuning, set $s \in [0, 1]$1 unless compute limitations force more aggressive pruning.

Empirical validation demonstrates that these guidelines track the true compute-loss envelope within 0.03–0.05 absolute loss, i.e., near-optimality across retrieval and semantic similarity benchmarks.

Compute Budget (FLOP)	Model Size ($s \in [0, 1]$2)	Data ($s \in [0, 1]$3, tokens)	Fine-tuning Method $s \in [0, 1]$4	LoRA Rank ($s \in [0, 1]$5)
$s \in [0, 1]$6–$s \in [0, 1]$7	largest possible	largest possible	full	—
$s \in [0, 1]$8–$s \in [0, 1]$9	up to limit	as above	LoRA	8–32

Note: Specific numerical boundaries may vary according to hardware and candidate pool.

6. Practical and Methodological Significance

The compute-optimal model design framework for embedding models formalizes the allocation of training resources, contrasts competing fine-tuning approaches, and provides a quantitative tool for architecture and data scaling under real-world constraints. This integrated methodology bridges empirical scaling laws and actionable design guidance, remedying the limitations of ad hoc model or data scaling. The phase transition between full and LoRA adaptation substantiates parameter-efficient fine-tuning as not only memory efficient, but provably optimal beyond certain compute thresholds (Ziarko et al., 2024).

The discretized search procedure and empirical phase diagram approach also anticipate the practical regime that embedding practitioners operate in: constrained resource pools and highly nonconvex choices over model, data, and adaptation protocol. The outlined recipe is robust over a range of LLM backbones and downstream embedding tasks.