Needle: Attention-Only Encoder-Decoder for On-Device Tool Calling
=================================================================

Target: NeurIPS 2026


## Core Claim

We distill a large cloud model's (Gemini) tool-calling capability into
a ~25M parameter attention-only encoder-decoder designed specifically
for this task. The key architectural finding: feedforward layers are
unnecessary for structured generation at small scale. Removing them
cuts 2/3 of per-layer parameters without losing quality, because tool
calling is retrieval-and-assembly, not reasoning, and attention alone
is sufficient for alignment and copying between input and output.


## What This Is

This is task-specific distillation with an architecture co-designed for
the distillation target. We are NOT claiming to train a model that
independently reasons about tools. We are claiming:

  1. Tool calling can be distilled into a much smaller model than
     previously demonstrated (25M vs 500M+ for existing approaches)

  2. The receiving architecture matters - an attention-only
     encoder-decoder absorbs this knowledge more efficiently than a
     standard transformer at the same parameter budget, because the
     task structure (alignment, copying, constrained output) maps
     directly onto what attention does

  3. The resulting model is small enough and fast enough for real
     on-device deployment with practical latency

The training data is synthesized by Gemini. The model's capability is
upper-bounded by Gemini's tool-calling ability. We are compressing that
capability into a form factor that can run on a phone.


## Why This Matters

On-device tool calling (voice assistants, smart home, wearables) needs
models that are small, fast, and accurate. Current approaches either:
  (a) Use large cloud models (accurate but slow, expensive, requires
      connectivity)
  (b) Distill into standard small transformers (still carry FFN
      parameters that don't help for this task)
  (c) Use heuristic/rule-based systems (brittle, can't generalize)

We show that a purpose-built architecture for distilling structured
generation achieves better compression than naive distillation into a
standard transformer, because the architecture matches the task.


## Paper Structure

### 1. Introduction

On-device tool calling is a constrained structured generation task.
The model sees a query and tool definitions, and must output valid JSON
tool calls. This is fundamentally different from open-ended text
generation:
  - The output vocabulary is constrained (tool names, parameter keys,
    JSON syntax)
  - The output content is extractive (argument values come from the
    query or are simple transforms of it)
  - The core operation is alignment (match query intent to tool, match
    query fragments to parameter slots)

These properties suggest that FFN - which provides per-position feature
transformation - may be unnecessary when the model is learning to
reproduce a teacher's structured outputs. Attention alone provides the
cross-position alignment and information routing that the task requires.

We validate this by distilling Gemini's tool-calling behavior into a
pure-attention encoder-decoder that removes FFN entirely and compensates
with gated residuals and task-specific heads. The architecture is
co-designed with the distillation target: structured generation from a
strong teacher into the smallest possible student.


### 2. Architecture

#### 2.1 Attention-Only Blocks

Encoder: x = x + sigmoid(gate) * SelfAttn(ZCRMSNorm(x))
Decoder: x = x + g_s * SelfAttn(ZCRMSNorm(x))
         x = x + g_c * CrossAttn(ZCRMSNorm(x))

No FFN. Per-sublayer learnable scalar gate initialized to 0 (sigmoid
= 0.5 at init). Model learns which layers matter.

Why gated residuals are necessary for attention-only: without FFN,
there is no per-position nonlinear rewriting. Standard residual forces
attention into a purely additive role. Gated residual lets the model
modulate per-layer contribution - can suppress unhelpful layers (g->0)
or amplify important ones (g->1).

#### 2.2 Encoder-Decoder

Bidirectional encoder for tool definitions (structured objects benefit
from seeing the full definition at once). Causal decoder for output
generation. Encoder output shared across three heads:

  Decoder:          generates tool calls
  Contrastive head: retrieves relevant tools from large catalogs
  Confidence head:  predicts whether to handle locally or route to cloud

#### 2.3 Contrastive Head (Tool Filtering)

Real deployments have 50-200 tools. They don't fit in 1024 encoder
tokens. Rather than adding a separate retrieval model, we reuse the
encoder: a lightweight MLP head projects mean-pooled encoder output
into a contrastive space for tool ranking.

  mean_pool -> Dense(d_model/4) -> ReLU -> Dense(128) -> L2-norm

CLIP-style training (0.1x weight, every 1000 steps). At inference,
encode the query and all candidate tools, rank by cosine similarity,
pass only top-k to the encoder for generation. No second model, no
extra memory, tool embeddings can be precomputed and cached.

#### 2.4 Confidence Head (Hybrid Routing)

MLP on mean-pooled encoder output:
  mean_pool -> Dense(d_model/4) -> ReLU -> Dense(1) -> sigmoid -> [0,1]

Trained post-hoc on decoder perplexity as self-supervised labels:
  confidence = sigmoid(-k * (log(ppl) - log(median_ppl)))

When confidence < threshold, skip decoder entirely and route to cloud.
Saves all decoder compute on queries the model can't handle.

#### 2.5 Constrained Decoding

Character-level trie-based logit masking during generation. JSON state
machine constrains tool names and argument keys to valid options from
the input. Guarantees structural validity with zero accuracy cost.


### 3. Training

Distillation from Gemini: synthetic (query, tools, answer) triples
with controlled diversity across 10 call types, 33 tool pools, and
10% synthesized tools. The dataset IS the distillation - we are
transferring Gemini's tool-calling behavior into the student.
See dataset.txt for the full pipeline.

Token-level loss weighting: base 1.0, names 2.0, keys 1.5, values 4.0.
Matches the error distribution (values > names > keys > structure).
See architecture.txt section 7.

Muon optimizer for attention projections (prevents representation
collapse in deep linear stacks), AdamW for everything else.
See architecture.txt section 6.


### 4. Evaluation Plan

This is the critical section. Without strong external evaluation the
paper is dead on arrival.

#### 4.1 External Benchmarks [MUST-HAVE]

  BFCL (Berkeley Function Calling Leaderboard):
    Standard benchmark. Scope mismatch (web APIs vs on-device) but
    necessary. If we do well despite never training on API-style tools,
    that's a generalization result. If we do poorly, we need to be
    honest about it and explain scope.

    Action items:
    - Adapt BFCL input format to encoder-decoder
    - Run all BFCL categories
    - Report alongside published results for other models

  ToolBench / API-Bank / ToolAlpaca:
    Secondary benchmarks. Less standardized but provide additional
    external validation.

#### 4.2 Head-to-Head vs Small LLMs [MUST-HAVE]

  Run small models on tool calling (either BFCL or a shared eval set):

  | Model            | Params | call_f1 | exact_match | tok/s | Memory |
  |------------------|--------|---------|-------------|-------|--------|
  | Needle (ours)    | ~25M   |         |             |       |        |
  | Qwen 2.5 0.5B   | 500M   |         |             |       |        |
  | Gemma 2 2B       | 2B     |         |             |       |        |
  | LLaMA 3.2 1B     | 1B     |         |             |       |        |
  | LLaMA 3.2 3B     | 3B     |         |             |       |        |
  | Phi-3-mini       | 3.8B   |         |             |       |        |

  We are 20-100x smaller. The story is NOT "we beat them" (we probably
  don't on general tool calling). The story is:
  - On in-scope tasks (on-device assistant), we are competitive despite
    being 20-100x smaller
  - We are 10-50x faster on edge hardware
  - We fit in <100MB memory vs 1-8GB for these models

  If we're within 10-15% of Qwen 0.5B at 1/20th the size, that's a
  strong result. If we're within 5% of Gemma 2B at 1/80th the size on
  on-device tasks, that's excellent.

#### 4.3 Latency on Real Hardware [MUST-HAVE]

  Without this, "on-device" is just a claim. Need actual numbers:

  | Device           | Model       | Latency (ms) | Memory (MB) |
  |------------------|-------------|--------------|-------------|
  | Pixel 8 (CPU)    | Needle      |              |             |
  | Pixel 8 (CPU)    | Qwen 0.5B   |              |             |
  | Raspberry Pi 5   | Needle      |              |             |
  | Apple M-series   | Needle      |              |             |

  This is the deployment story. If Needle runs in <50ms on a phone CPU
  while Qwen 0.5B takes >500ms, that's the practical contribution.

  Action items:
  - Export model to ONNX or TFLite
  - Benchmark on 2-3 real devices
  - Include memory footprint

#### 4.4 Ablations on Internal Val Set [MUST-HAVE]

  These use our synthetic val set but that's OK for ablations -
  we're comparing our own variants against each other, not claiming
  absolute performance.

  A. FFN vs No-FFN (THE critical ablation):
     Same parameter budget, same data, same recipe.
     | Condition  | Params | call_f1 | exact_match | tok/s |
     Must show no-FFN is not worse. If it's better, even stronger.

  B. Gated vs Standard vs No Residual:
     | Condition  | Converges? | Final PPL | call_f1 |
     Must show gated is necessary for attention-only.

  C. Encoder-Decoder vs Decoder-Only:
     Same budget, same data.
     | Architecture | call_f1 | exact_match |
     Must show enc-dec wins on tool calling.

  D. With vs Without Contrastive:
     lambda=0.1 vs lambda=0, generation metrics only.
     | Condition | name_f1 | call_f1 |
     Shows contrastive helps generation, not just retrieval.

  E. Constrained vs Unconstrained Decoding:
     | Condition | name_f1 | key_acc | call_f1 |
     Shows free accuracy from structural constraints.

#### 4.5 Analysis [STRONG]

  A. Gate value analysis: which layers matter? Do early/late/middle
     layers behave differently? Are any gates near 0 (dormant)?

  B. Confidence calibration: reliability diagram (predicted vs actual
     accuracy). Routing accuracy curve at various thresholds.

  C. Retrieval quality: Recall@k, MRR on val set. Stratified by tool
     count.

  D. Depth vs width: 6+6 d=768 vs 12+12 d=512 vs 24+24 d=256 at
     fixed budget. Informs scaling for attention-only.


### 5. Related Work

  Knowledge distillation:
    Hinton et al. 2015, DistilBERT, TinyBERT. We distill from a
    generative teacher (Gemini) into a task-specific student, not a
    general-purpose one. The novelty is not distillation itself but
    the receiving architecture being co-designed for the task.

  Tool calling models:
    Gorilla, ToolLLaMA, NexusRaven, Toolformer - all decoder-only,
    all >1B params. We use encoder-decoder at 25M. These are trained
    on human-curated or self-generated data; we use teacher distillation
    explicitly.

  Attention-only / FFN-free:
    Limited prior work. Some theoretical results on attention
    expressiveness (Yun et al. 2020 - transformers are universal
    approximators, but the proof uses FFN). We show empirically that
    for structured generation distillation, attention alone suffices.

  Hybrid inference / routing:
    Speculative decoding, mixture-of-experts routing. Our confidence
    head is simpler - binary local/cloud decision based on encoder
    output. The student knows when it can't reproduce the teacher.

  Constrained decoding:
    PICARD (SQL), grammar-constrained generation (Scholak et al.).
    We apply character-level trie masking to tool-call JSON.


### 6. Limitations

  - This is distillation, not independent learning. The model's
    capability is upper-bounded by Gemini's tool-calling ability.
    If Gemini gets a tool call wrong in training data, the student
    learns that error.
  - Evaluated primarily on on-device assistant tasks. May not
    generalize to complex multi-step API orchestration.
  - The no-FFN finding may be specific to structured generation.
    Open-ended text generation likely still needs FFN. We claim
    task-specific sufficiency, not general sufficiency.
  - Confidence head calibration depends on training data distribution.
    Out-of-distribution queries may have poorly calibrated confidence.
  - Small vocab (8192) limits applicability beyond tool calling.
  - The architecture is purpose-built. It cannot be repurposed for
    general text generation, summarization, or reasoning tasks without
    fundamental changes (re-adding FFN at minimum).


## Reviewer Objections and Responses

  Q: "Evaluation is only on synthetic data."
  A: We evaluate on BFCL, head-to-head vs small LLMs, and on real
     hardware. The synthetic val set is used only for ablations
     (comparing our own variants). Absolute performance claims are
     backed by external benchmarks.

  Q: "This is just distillation from Gemini."
  A: Yes, it is. We say so explicitly. The contribution is not the
     distillation itself but the receiving architecture. Given the SAME
     teacher data, an attention-only encoder-decoder absorbs tool-calling
     knowledge more efficiently than a standard transformer (ablation A).
     The question isn't "can you distill?" - it's "what's the best
     student architecture for structured generation distillation?"

  Q: "No-FFN only works because the task is easy."
  A: Tool calling is not easy - BFCL leaderboard shows even large
     models make errors. It is STRUCTURED, which is different from
     easy. The claim is that structure enables attention-only, not
     that the task is trivial.

  Q: "The model is too small to be useful."
  A: [Show latency numbers on phone.] This runs in <Xms on a phone
     CPU. No other tool-calling model can do that. Size is the point.

  Q: "Too many components, no single deep contribution."
  A: The contribution is the co-design. We make three specific claims:
     (1) FFN is unnecessary for structured generation at small scale
     (ablation A), (2) shared contrastive training improves generation
     (ablation D), (3) encoder-based confidence routing saves compute
     without losing accuracy (Sec 4.5B). Everything else is recipe.

  Q: "The no-FFN argument is hand-wavy."
  A: We make an empirical claim, not a theoretical one. We do not
     prove attention alone is sufficient in general. We show it is
     sufficient for this task at this scale. The ablation is the proof.


## Action Items

Priority order:

  1. Run the FFN ablation. This is the paper.
     If no-FFN loses badly, there is no paper.

  2. Run BFCL. Need external benchmark numbers.
     Adapt input format, run all categories, report honestly.

  3. Run Qwen 0.5B / Gemma 2B / LLaMA 3.2 1B on tool calling.
     Need the comparison table.

  4. Get latency on real hardware.
     Export to ONNX/TFLite. Benchmark on phone + RPi.

  5. Run remaining ablations (gated residual, enc-dec vs dec-only,
     contrastive, constrained).

  6. Gate value analysis and confidence calibration plots.

  7. Write the paper.

Items 1-4 are blocking. Without them, no submission.
Items 5-6 strengthen but don't make or break.
