Llama3_8B_Ambiref
meta-llama/Llama-3.1-8B (base, not instruct), fine-tuned with LoRA on the boxes entity-tracking task (AmbiRef split from Kim & Schuster, 2023). The LoRA adapter has been merged into the base weights, so this repo loads as a standard LlamaForCausalLM.
The model reads a description of boxes and their contents, then a sequence of operations (put / remove / move), and completes the final contents of a queried box.
Usage
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
repo = "mousse95/Llama3_8B_Ambiref"
tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(repo, torch_dtype=torch.bfloat16, device_map="auto")
prompt = (
"Box 0 contains the big rose, Box 1 contains the green plate, Box 2 contains the blue camera. "
"Move the plate from Box 1 to Box 0. Remove the camera from Box 2. "
"Box 0 contains <extra_id_0> ."
)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=32, do_sample=False)
print(tokenizer.decode(out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True))
Prompt format
Prompts end with Box N contains <extra_id_0> . The target completion is a space, then <extra_id_0>, then the box contents joined by and (or nothing), then EOS. <extra_id_0> is plain text here, not a special token. It is kept from the original T5-style data.
Training
Trained with train_base_8b_entitytracking_lora_fixed.py, which is included in this repo. The full run record is in training/training_summary.json.
Data. 19,500 training and 500 dev examples, sampled without replacement as follows:
| Bucket | Share | Train examples |
|---|---|---|
| Easy (circuit-discovery entity-tracking set) | 10% | 1,950 |
NumOps = 2 (AmbiRef train-t5.jsonl) | 60% | 11,700 |
NumOps ≥ 3 (AmbiRef train-t5.jsonl) | 30% | 5,850 |
Leakage controls: every row from one source scene goes to exactly one split (train or dev), no scene/operation chain appears in both train and dev, and exact-pair, prompt and scene overlap with the AmbiRef test-t5.jsonl was checked and removed (0 overlaps found).
Objective. Completion-only causal LM loss: prompt tokens are masked (-100), and only the answer and EOS are supervised.
Hyperparameters.
| LoRA | r = 64, α = 128, dropout = 0.05, target = all-linear, RS-LoRA |
| Trainable params | 167.8M (2.05% of 8.2B) |
| Epochs | 3 |
| Effective batch size | 64 (16 × 4 gradient accumulation) |
| Learning rate | 1e-4, warmup ratio 0.05, weight decay 0.01 |
| Optimizer | adamw_torch_fused, max grad norm 1.0 |
| Precision / max length | bf16 / 512 tokens |
| Seed | 42 |
| Hardware | 1× NVIDIA B300, ~16 min |
Results at the end of training.
| Metric | Value |
|---|---|
| Train loss | 0.058 |
| Dev loss | 0.034 |
| Greedy exact match, 64 train examples | 1.00 |
| Greedy exact match, 64 dev examples | 0.83 |
The exact-match figures come from a small sanity check run during training, not from a full evaluation.
Data
The base data is the AmbiRef split (boxes-dataset-v1/few_shot_boxes_nso_exp2_max3_ambiref) of the boxes entity-tracking dataset.
The authors distribute it as a password-protected ZIP to keep it out of LM training data, so it is not republished here. Download it from their repository. The additional circuit / mapping / student-head splits used in the downstream experiments were derived from it for this project and are not public either.
License
Built with Llama. This model is a derivative of Llama 3.1 and is distributed under the Llama 3.1 Community License. Use of the model is also subject to the Llama 3.1 Acceptable Use Policy.