SAFT: Structure-Aware Fine-Tuning of Large Language Models for AMR-to-Text Generation

This repository contains the code for SAFT (Structure-Aware Fine-Tuning), a method for fine-tuning causal language models on AMR graph–text pairs using LoRA and optional graph positional encodings derived from the magnetic Laplacian. SAFT injects graph structure into pretrained LLMs without modifying their architecture and achieves strong results on AMR-to-text generation.

The method is described in:

SAFT: Structure-Aware Fine-Tuning of Large Language Models for AMR-to-Text Generation
Transactions on Machine Learning Research (TMLR), 2026
[OpenReview]

Part of the experiments use a modified litgpt codebase: https://github.com/RafiqKamel/litgpt/tree/graphPE.

---

Overview

SAFT fine-tunes causal LMs on AMR graph–text pairs. For each AMR graph, we compute direction-sensitive positional encodings from its magnetic Laplacian and project them into the model’s embedding space. During fine-tuning, these graph encodings are added to token embeddings aligned with AMR nodes, providing a structure-aware inductive bias while keeping the base model and attention layers unchanged.

The code supports:

• Fine-tuning with and without graph positional encodings.
• LoRA-based adaptation of causal LMs.
• Sentence-level and document-level AMR graphs.
• Evaluation with BLEU and chrF by default, with optional METEOR, BERTScore, COMET, and BLEURT.

---

Requirements

• Linux (tested) with recent CUDA-capable GPU(s).
• Python and CUDA versions compatible with PyTorch in the Pixi environment.
• Access to AMR 3.0 (LDC2020T02) under your own LDC license.
• Sufficient disk space for processed AMR data and eigen-decompositions.

We manage all dependencies with Pixi, a reproducible, lockfile-based environment manager built on top of conda-forge.

---

Installation

1. Install Pixi (once per machine). See the official instructions at https://pixi.sh/, for example:

   curl -fsSL https://pixi.sh/install.sh | bash

   Then restart your shell so pixi is on your PATH.

2. Create the SAFT environment in this repository:

   # from the repo root
   pixi install

3. Verify the install by running the test suite:

   pixi run test

The default Pixi environment uses CUDA-enabled PyTorch and PyTorch’s scaled dot-product attention (SDPA). FlashAttention is not bundled by default; see below for enabling it.

---

Data: AMR 3.0 and preprocessing

We do not distribute AMR data, checkpoints, or processed artifacts. To run SAFT you must obtain AMR 3.0 (LDC2020T02) under your own LDC agreement and place the raw data in a directory layout similar to:

/path/to/amr/
  train/
  dev/
  test/
  # or per-split directories as provided by LDC

Set environment variables pointing to your processed data and runs root:

export SAFT_PROCESSED_DIR=/path/to/processed-amr
export SAFT_RUNS_ROOT=/path/to/saft-runs

Then build the HF-compatible dataset:

pixi run python tools/data/build_dataset.py \
  --raw-dir /path/to/amr/split \
  --output-dir "$SAFT_PROCESSED_DIR/hf"

This step converts AMR 3.0 into tokenized graph–text pairs and writes them under $SAFT_PROCESSED_DIR/hf. See data/README.md for details on the expected directory structure and split conventions.

Next, compute graph positional encodings (e.g., magnetic Laplacian eigenvectors) for all graphs:

pixi run build-eig

This will read from $SAFT_PROCESSED_DIR/hf, perform eigen-decompositions per graph, and store the resulting encodings alongside the processed dataset.

For more detailed setup instructions and options (e.g., document-level graphs, custom splits), see docs/setup.md.

---

Quickstart: fine-tuning and evaluation

Once the environment and data are prepared, you can run a minimal smoke test fine-tuning and evaluation.

Fine-tuning

Run a small “smoke” experiment to check that everything works end-to-end:

pixi run finetune --config-name smoke

This will:

• Load the configuration smoke (see the config directory for details).
• Fine-tune a small model on a reduced subset of AMR 3.0.
• Write logs, checkpoints, and metadata under:

$SAFT_RUNS_ROOT/smoke/smoke-run/

You can adjust the configuration (model size, batch size, learning rate, LoRA rank, graph encodings on/off, etc.) by editing or extending the corresponding config file.

Evaluation

Evaluate a trained run on the validation split:

pixi run eval "$SAFT_RUNS_ROOT/smoke/smoke-run" \
  --split val \
  --max-samples 64

Core evaluation reports BLEU and chrF. By default we use standard implementations for these metrics. Additional metrics are available as optional integrations:

• METEOR
• BERTScore
• COMET
• BLEURT

See the setup guide in docs/setup.md for how to install the corresponding packages and enable each metric in the evaluation configuration.

---

FlashAttention (optional)

The default configuration uses PyTorch’s built-in scaled dot-product attention implementation, which works out-of-the-box.

To enable FlashAttention-2, set:

model:
  attn_implementation: flash_attention_2

in your configuration and ensure that a compatible flash-attn build is installed for your hardware and CUDA/PyTorch stack. We intentionally do not bundle FlashAttention in the Pixi environment, since it has strict GPU and CUDA requirements and may not build on all systems.

Only enable FlashAttention after you have installed and tested it separately.

---

Development and quality checks

We provide several tasks to keep the codebase consistent:

# Run linters (e.g., ruff, flake8)
pixi run lint

# Check formatting (e.g., black, isort) without modifying files
pixi run fmt-check

# Run static type checks (e.g., mypy, pyright)
pixi run typecheck

# Verify that all configuration files are valid and self-consistent
pixi run verify-configs

These commands are intended for contributors and for ensuring reproducible experiments.

---

License

Most of this repository is licensed under the MIT License; see LICENSE.

The AMR linearization module derived from SPRING is licensed under CC BY-NC-SA 4.0; see NOTICE for details. This means that commercial use of that component is restricted by the original SPRING license.

---

Citation

If you build on this work, please consider citing it:

@article{kamel2026saft,
    title={{SAFT}: Structure-Aware Fine-Tuning of Large Language Models for {AMR}-to-Text Generation},
    author={Anonymous},
    journal={Transactions on Machine Learning Research},
    year={2026},
    url={https://openreview.net/forum?id=QZoUMyzYDB}
}